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 2a8d88a ivfpq: Reuse query distance tables across IVF lists (#67)
2a8d88a is described below
commit 2a8d88ac67b40944c9728195a3ecbc298afe06cc
Author: jerry <[email protected]>
AuthorDate: Mon Aug 3 21:40:15 2026 +0800
ivfpq: Reuse query distance tables across IVF lists (#67)
---
c/test_vindex.c | 12 +
core/Cargo.toml | 4 +
core/benches/ivfpq_batch_reuse_bench.rs | 165 ++++++++
core/src/index.rs | 39 +-
core/src/io.rs | 6 +-
core/src/ivfpq.rs | 454 +++++++++++++++++++--
cpp/test_vindex.cpp | 2 +
ffi/src/lib.rs | 137 +++++++
include/paimon_vindex.hpp | 21 +-
.../paimon/index/vector/VectorSearchParams.java | 47 ++-
.../index/vector/VectorIndexJavaApiTest.java | 13 +
jni/src/lib.rs | 28 ++
python/paimon_vindex/__init__.py | 52 ++-
python/paimon_vindex/_ffi.py | 34 ++
python/tests/test_vindex.py | 22 +
15 files changed, 980 insertions(+), 56 deletions(-)
diff --git a/c/test_vindex.c b/c/test_vindex.c
index fbd1c93..ac8fd50 100644
--- a/c/test_vindex.c
+++ b/c/test_vindex.c
@@ -339,6 +339,18 @@ static void run_roundtrip(
}
assert_id_in_cluster(batch_ids[0], 0);
assert_id_in_cluster(batch_ids[1], 1);
+ struct PaimonVindexSearchParamsV2 batch_params_v2 = {
+ .top_k = 1,
+ .search_width = batch_params.search_width,
+ .width = batch_params.width,
+ .ivfpq_batch_table_reuse = PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_OFF,
+ .ivfpq_batch_table_reuse_max_bytes = 1};
+ if (paimon_vindex_reader_search_batch_v2(
+ reader, queries, 2, batch_params_v2, batch_ids, batch_distances,
2) != 0) {
+ fail_ffi("reader search batch v2 failed");
+ }
+ assert_id_in_cluster(batch_ids[0], 0);
+ assert_id_in_cluster(batch_ids[1], 1);
paimon_vindex_reader_free(reader);
free(buf.data);
free(data);
diff --git a/core/Cargo.toml b/core/Cargo.toml
index 4b3c126..317afa5 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -53,3 +53,7 @@ harness = false
[[bench]]
name = "diskann_bench"
harness = false
+
+[[bench]]
+name = "ivfpq_batch_reuse_bench"
+harness = false
diff --git a/core/benches/ivfpq_batch_reuse_bench.rs
b/core/benches/ivfpq_batch_reuse_bench.rs
new file mode 100644
index 0000000..1445015
--- /dev/null
+++ b/core/benches/ivfpq_batch_reuse_bench.rs
@@ -0,0 +1,165 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use paimon_vindex_core::distance::MetricType;
+use paimon_vindex_core::io::{write_index, IVFPQIndexReader, PosWriter};
+use paimon_vindex_core::ivfpq::{
+ search_batch_reader_with_reuse_mode, IVFPQIndex, IvfPqBatchTableReuseMode,
+};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use std::hint::black_box;
+use std::io::Cursor;
+use std::time::Instant;
+
+const D: usize = 768;
+const M: usize = 16;
+const NLIST: usize = 256;
+const NPROBE: usize = 8;
+const NQ: usize = 256;
+const K: usize = 10;
+const ROWS_PER_LIST: usize = 390;
+const ROUNDS: usize = 100;
+
+#[derive(Clone, Copy)]
+struct CpuTimes {
+ user: f64,
+ system: f64,
+}
+
+#[derive(Clone, Copy)]
+struct Sample {
+ wall: f64,
+ user: f64,
+ system: f64,
+}
+
+#[cfg(unix)]
+fn cpu_times() -> CpuTimes {
+ let mut usage = std::mem::MaybeUninit::<libc::rusage>::uninit();
+ let result = unsafe { libc::getrusage(libc::RUSAGE_SELF,
usage.as_mut_ptr()) };
+ assert_eq!(result, 0);
+ let usage = unsafe { usage.assume_init() };
+ CpuTimes {
+ user: usage.ru_utime.tv_sec as f64 + usage.ru_utime.tv_usec as f64 /
1_000_000.0,
+ system: usage.ru_stime.tv_sec as f64 + usage.ru_stime.tv_usec as f64 /
1_000_000.0,
+ }
+}
+
+#[cfg(not(unix))]
+fn cpu_times() -> CpuTimes {
+ CpuTimes {
+ user: 0.0,
+ system: 0.0,
+ }
+}
+
+fn search(
+ reader: &mut IVFPQIndexReader<Cursor<Vec<u8>>>,
+ queries: &[f32],
+ mode: IvfPqBatchTableReuseMode,
+) -> ((Vec<i64>, Vec<f32>), Sample) {
+ let cpu_before = cpu_times();
+ let started = Instant::now();
+ let result = search_batch_reader_with_reuse_mode(reader, queries, NQ, K,
NPROBE, mode).unwrap();
+ let wall = started.elapsed().as_secs_f64();
+ let cpu_after = cpu_times();
+ black_box(result.0.iter().fold(0i64, |sum, id| sum.wrapping_add(*id)));
+ (
+ result,
+ Sample {
+ wall,
+ user: cpu_after.user - cpu_before.user,
+ system: cpu_after.system - cpu_before.system,
+ },
+ )
+}
+
+fn percentile(samples: &[Sample], percentile: usize, value: impl Fn(&Sample)
-> f64) -> f64 {
+ let mut values = samples.iter().map(value).collect::<Vec<_>>();
+ values.sort_by(f64::total_cmp);
+ values[(percentile * values.len()).div_ceil(100).saturating_sub(1)]
+}
+
+fn main() {
+ let mut rng = StdRng::seed_from_u64(42);
+ let mut index = IVFPQIndex::new(D, NLIST, M, MetricType::InnerProduct,
false);
+ index.quantizer_centroids = (0..NLIST * D)
+ .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();
+ 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)
+ .map(|id| id as i64)
+ .collect();
+ index.codes[list_id] = (0..ROWS_PER_LIST * M).map(|_|
rng.gen()).collect();
+ }
+ let queries = (0..NQ * D)
+ .map(|_| rng.gen_range(-1.0f32..1.0))
+ .collect::<Vec<_>>();
+ let mut bytes = Vec::new();
+ write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+
+ let modes = [
+ IvfPqBatchTableReuseMode::Off,
+ IvfPqBatchTableReuseMode::On,
+ IvfPqBatchTableReuseMode::Auto,
+ ];
+ let mut readers = [
+ IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(),
+ IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(),
+ IVFPQIndexReader::open(Cursor::new(bytes)).unwrap(),
+ ];
+ let expected = search(&mut readers[0], &queries, modes[0]).0;
+ assert_eq!(search(&mut readers[1], &queries, modes[1]).0, expected);
+ assert_eq!(search(&mut readers[2], &queries, modes[2]).0, expected);
+
+ let mut samples = [Vec::new(), Vec::new(), Vec::new()];
+ let orders = [[0, 1, 2], [1, 2, 0], [2, 0, 1]];
+ for round in 0..ROUNDS {
+ for &mode_index in &orders[round % orders.len()] {
+ let (_, sample) = search(&mut readers[mode_index], &queries,
modes[mode_index]);
+ samples[mode_index].push(sample);
+ }
+ }
+
+ println!(
+ "shape: d={D} m={M} nlist={NLIST} nprobe={NPROBE} nq={NQ} vectors={}
threads={} rounds={ROUNDS}",
+ NLIST * ROWS_PER_LIST,
+ rayon::current_num_threads()
+ );
+
println!("mode,percentile,wall_ms,user_cpu_ms,sys_cpu_ms,total_cpu_ms,cpu/wall");
+ for (mode, samples) in ["off", "on", "auto"].into_iter().zip(&samples) {
+ for p in [50, 90, 95, 99] {
+ let wall = percentile(samples, p, |sample| sample.wall);
+ let user = percentile(samples, p, |sample| sample.user);
+ let system = percentile(samples, p, |sample| sample.system);
+ let total = percentile(samples, p, |sample| sample.user +
sample.system);
+ println!(
+ "{mode},p{p},{:.3},{:.3},{:.3},{:.3},{:.2}",
+ wall * 1_000.0,
+ user * 1_000.0,
+ system * 1_000.0,
+ total * 1_000.0,
+ total / wall
+ );
+ }
+ }
+}
diff --git a/core/src/index.rs b/core/src/index.rs
index 41d16bc..cc2ed7c 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -33,11 +33,12 @@ use crate::ivfflat_io::{
search_batch_ivfflat_reader, search_batch_ivfflat_reader_roaring_filter,
write_ivfflat_index,
IVFFlatIndexReader, IVFFLAT_MAGIC,
};
-pub use crate::ivfpq::IvfPqBatchTableReuseMode;
use crate::ivfpq::{
- search_batch_reader_roaring_filter_with_reuse_mode,
search_batch_reader_with_reuse_mode,
- search_with_reader, search_with_reader_roaring_filter, IVFPQIndex,
+ search_batch_reader_roaring_filter_with_reuse_mode_and_budget,
+ search_batch_reader_with_reuse_mode_and_budget, search_with_reader,
+ search_with_reader_roaring_filter, IVFPQIndex,
};
+pub use crate::ivfpq::{IvfPqBatchTableReuseMode,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES};
use crate::ivfrq::IVFRQIndex;
pub use crate::ivfrq_io::IVFRQSearchStats;
use crate::ivfrq_io::{
@@ -991,6 +992,7 @@ pub struct VectorSearchParams {
pub search_width: SearchWidth,
pub width: usize,
pub ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode,
+ pub ivfpq_batch_table_reuse_max_bytes: usize,
}
impl VectorSearchParams {
@@ -1000,6 +1002,7 @@ impl VectorSearchParams {
search_width: SearchWidth::IvfNProbe,
width: nprobe,
ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto,
+ ivfpq_batch_table_reuse_max_bytes:
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
}
}
@@ -1009,6 +1012,7 @@ impl VectorSearchParams {
search_width: SearchWidth::DiskAnnLSearch,
width: l_search,
ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto,
+ ivfpq_batch_table_reuse_max_bytes:
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
}
}
@@ -1018,6 +1022,7 @@ impl VectorSearchParams {
search_width: SearchWidth::Auto,
width: 0,
ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto,
+ ivfpq_batch_table_reuse_max_bytes:
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
}
}
@@ -1026,6 +1031,11 @@ impl VectorSearchParams {
self
}
+ pub fn with_ivfpq_batch_table_reuse_max_bytes(mut self, max_bytes: usize)
-> Self {
+ self.ivfpq_batch_table_reuse_max_bytes = max_bytes;
+ self
+ }
+
pub fn configured_ivf_nprobe(self) -> Option<usize> {
(self.search_width == SearchWidth::IvfNProbe).then_some(self.width)
}
@@ -1035,7 +1045,11 @@ impl VectorSearchParams {
}
fn validate(self) -> io::Result<()> {
- validate_positive(self.top_k, "top_k")
+ validate_positive(self.top_k, "top_k")?;
+ validate_positive(
+ self.ivfpq_batch_table_reuse_max_bytes,
+ "IVF-PQ batch table reuse max bytes",
+ )
}
fn resolve_ivf_nprobe(
@@ -1773,13 +1787,14 @@ impl<R: SeekRead> VectorIndexReader<R> {
params.top_k,
total_vectors,
|nprobe| {
- search_batch_reader_with_reuse_mode(
+ search_batch_reader_with_reuse_mode_and_budget(
reader,
queries,
query_count,
params.top_k,
nprobe,
params.ivfpq_batch_table_reuse,
+ params.ivfpq_batch_table_reuse_max_bytes,
)
},
)
@@ -1890,7 +1905,7 @@ impl<R: SeekRead> VectorIndexReader<R> {
params.top_k,
matching_count.unwrap_or(total_vectors),
|nprobe| {
- search_batch_reader_roaring_filter_with_reuse_mode(
+
search_batch_reader_roaring_filter_with_reuse_mode_and_budget(
reader,
queries,
query_count,
@@ -1898,6 +1913,7 @@ impl<R: SeekRead> VectorIndexReader<R> {
nprobe,
roaring_filter_bytes,
params.ivfpq_batch_table_reuse,
+ params.ivfpq_batch_table_reuse_max_bytes,
)
},
)
@@ -2925,12 +2941,23 @@ mod tests {
params.ivfpq_batch_table_reuse,
IvfPqBatchTableReuseMode::Auto
);
+ assert_eq!(params.ivfpq_batch_table_reuse_max_bytes, 512 * 1024 *
1024);
assert_eq!(
params
.with_ivfpq_batch_table_reuse(IvfPqBatchTableReuseMode::Off)
.ivfpq_batch_table_reuse,
IvfPqBatchTableReuseMode::Off
);
+ assert_eq!(
+ params
+ .with_ivfpq_batch_table_reuse_max_bytes(128 * 1024 * 1024)
+ .ivfpq_batch_table_reuse_max_bytes,
+ 128 * 1024 * 1024
+ );
+ assert!(params
+ .with_ivfpq_batch_table_reuse_max_bytes(0)
+ .validate()
+ .is_err());
}
#[test]
diff --git a/core/src/io.rs b/core/src/io.rs
index f66eade..87b7dd3 100644
--- a/core/src/io.rs
+++ b/core/src/io.rs
@@ -893,7 +893,7 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
pub(crate) fn for_each_streamed_list_chunk(
&mut self,
list_id: usize,
- mut consume: impl FnMut(&[i64], &[u8]),
+ mut consume: impl FnMut(&ProductQuantizer, &[i64], &[u8]),
) -> io::Result<()> {
self.ensure_loaded()?;
let count = self.list_counts[list_id] as usize;
@@ -964,7 +964,7 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
.pread(&mut [ReadRequest::new(chunk_offset,
payload.codes_mut())])?;
}
let row_end = row_start + chunk_rows;
- consume(&ids[row_start..row_end], payload.codes());
+ consume(&self.pq, &ids[row_start..row_end], payload.codes());
row_start = row_end;
}
Ok(())
@@ -1349,7 +1349,7 @@ mod tests {
let mut actual_ids = Vec::new();
let mut actual_codes = Vec::new();
streamed_reader
- .for_each_streamed_list_chunk(0, |ids, codes| {
+ .for_each_streamed_list_chunk(0, |_, ids, codes| {
actual_ids.extend_from_slice(ids);
actual_codes.extend_from_slice(codes);
})
diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs
index f0c35ad..75afcf1 100644
--- a/core/src/ivfpq.rs
+++ b/core/src/ivfpq.rs
@@ -19,7 +19,7 @@ use crate::distance::{
fvec_inner_product, fvec_madd, fvec_normalize, pq_distance_four_codes,
pq_distance_from_table,
MetricType,
};
-use crate::index_io_util::{ivf_payload_is_oversized, MAX_IVF_BATCH_READ_BYTES};
+use crate::index_io_util::ivf_payload_is_oversized;
use crate::io::{IVFPQIndexReader, InvertedListPayload, SeekRead};
use crate::kmeans::{self, KMeansConfig};
use crate::opq::OPQMatrix;
@@ -29,6 +29,7 @@ use roaring::RoaringTreemap;
use std::borrow::Cow;
use std::collections::HashSet;
use std::io;
+use std::sync::OnceLock;
pub trait RowIdFilter: Sync {
fn contains(&self, id: i64) -> bool;
@@ -402,8 +403,8 @@ impl IVFPQIndex {
let coarse_dists = &all_coarse_dists[qi];
let mut heap = TopKHeap::new(k);
- let mut sim_table = vec![0.0f32; m * ksub];
-
+ let mut sim_table = Vec::new();
+ let mut non_residual_table_ready = false;
let ip_table = if use_precomputed {
let mut t = vec![0.0f32; m * ksub];
self.pq.compute_inner_product_table(query, &mut t);
@@ -422,6 +423,15 @@ impl IVFPQIndex {
continue;
}
+ if sim_table.is_empty() {
+ sim_table.resize(m * ksub, 0.0);
+ }
+ if !self.by_residual && !non_residual_table_ready {
+ self.pq
+ .compute_distance_table(query, self.metric, &mut
sim_table);
+ non_residual_table_ready = true;
+ }
+
// Precomputed sim_table omits ||q-c||²; add it as dis0.
// Non-precomputed path computes from residual_query,
already full distance.
let dis0 = if use_precomputed {
@@ -438,7 +448,7 @@ impl IVFPQIndex {
-2.0,
&mut sim_table,
);
- } else {
+ } else if self.by_residual {
self.compute_list_table(query, list_id, &mut
sim_table);
}
@@ -924,6 +934,7 @@ fn has_matching_rows(matching_rows: Option<&MatchingRows>)
-> bool {
// Below this size, table construction and Rayon scheduling dominate the saved
// per-query/list distance-table work.
const MIN_EPHEMERAL_PRECOMPUTE_QUERIES: usize = 64;
+pub const DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES: usize = 512 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
@@ -948,6 +959,7 @@ fn ephemeral_precomputed_table_fits_budget(
query_scratch_count: usize,
m: usize,
ksub: usize,
+ max_bytes: usize,
) -> bool {
if matching_list_count == 0 {
return false;
@@ -958,7 +970,7 @@ fn ephemeral_precomputed_table_fits_budget(
.and_then(|tables| tables.checked_mul(m))
.and_then(|values| values.checked_mul(ksub))
.and_then(|values| values.checked_mul(std::mem::size_of::<f64>()))
- .is_some_and(|bytes| bytes <= MAX_IVF_BATCH_READ_BYTES)
+ .is_some_and(|bytes| bytes <= max_bytes)
}
#[cfg(test)]
@@ -1312,6 +1324,9 @@ struct ReaderSearchContext<'a> {
q: &'a [f32],
ip_table: &'a [f32],
use_precomputed: bool,
+ shared_sim_table: Option<&'a OnceLock<Vec<f32>>>,
+ #[cfg(test)]
+ distance_table_builds: Option<&'a std::sync::atomic::AtomicUsize>,
d: usize,
m: usize,
ksub: usize,
@@ -1401,6 +1416,7 @@ pub fn search_with_reader_filter<R: SeekRead>(
} else {
Vec::new()
};
+ let shared_sim_table = OnceLock::new();
let mut heap = TopKHeap::new(k);
@@ -1428,14 +1444,25 @@ pub fn search_with_reader_filter<R: SeekRead>(
let first_list = read_list_ids[batch_start];
if ivf_payload_is_oversized(reader.list_payload_len(first_list)?) {
let (_, _, dis0) = lists_to_read[batch_start];
- let sim_table = reader_sim_table(reader, first_list, &q,
&ip_table, use_precomputed);
+ let sim_table = by_residual
+ .then(|| reader_sim_table(reader, first_list, &q, &ip_table,
use_precomputed));
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, |ids, codes| {
+ reader.for_each_streamed_list_chunk(first_list, |pq, ids, codes| {
let positions = matching_rows(ids, filter);
+ if positions.as_ref().is_some_and(MatchingRows::is_empty) {
+ return;
+ }
+ let sim_table = sim_table.as_deref().unwrap_or_else(|| {
+ shared_sim_table.get_or_init(|| {
+ let mut table = vec![0.0f32; m * ksub];
+ pq.compute_distance_table(&q, metric, &mut table);
+ table
+ })
+ });
scan_reader_codes(
- &sim_table,
+ sim_table,
codes,
ids,
m,
@@ -1478,6 +1505,9 @@ pub fn search_with_reader_filter<R: SeekRead>(
q: &q,
ip_table: &ip_table,
use_precomputed,
+ shared_sim_table: (!by_residual).then_some(&shared_sim_table),
+ #[cfg(test)]
+ distance_table_builds: None,
d,
m,
ksub,
@@ -1543,9 +1573,18 @@ fn scan_reader_list(
if matching_rows.is_some_and(MatchingRows::is_empty) {
return;
}
- fill_reader_sim_table(entry.list_id, ctx, &mut scratch.sim_table);
+ let sim_table = if let Some(table) = ctx.shared_sim_table {
+ table.get_or_init(|| {
+ let mut sim_table = Vec::new();
+ fill_reader_sim_table(entry.list_id, ctx, &mut sim_table);
+ sim_table
+ })
+ } else {
+ fill_reader_sim_table(entry.list_id, ctx, &mut scratch.sim_table);
+ &scratch.sim_table
+ };
scan_reader_codes(
- &scratch.sim_table,
+ sim_table,
entry.codes(),
&entry.ids,
ctx.m,
@@ -1564,6 +1603,12 @@ fn fill_reader_sim_table(list_id: usize, ctx:
&ReaderSearchContext<'_>, sim_tabl
let m = ctx.m;
let ksub = ctx.ksub;
sim_table.resize(m * ksub, 0.0);
+ #[cfg(test)]
+ if !ctx.use_precomputed {
+ if let Some(builds) = ctx.distance_table_builds {
+ builds.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+ }
+ }
if ctx.use_precomputed {
let tab_base = list_id * m * ksub;
fvec_madd(
@@ -1598,6 +1643,9 @@ fn reader_sim_table<R: SeekRead>(
q: query,
ip_table,
use_precomputed,
+ shared_sim_table: None,
+ #[cfg(test)]
+ distance_table_builds: None,
d: reader.d,
m: reader.m,
ksub: reader.ksub,
@@ -1701,7 +1749,36 @@ pub fn search_batch_reader_with_reuse_mode<R: SeekRead>(
nprobe: usize,
reuse_mode: IvfPqBatchTableReuseMode,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
- search_batch_reader_filter_with_reuse_mode(reader, queries, nq, k, nprobe,
None, reuse_mode)
+ search_batch_reader_with_reuse_mode_and_budget(
+ reader,
+ queries,
+ nq,
+ k,
+ nprobe,
+ reuse_mode,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ )
+}
+
+pub fn search_batch_reader_with_reuse_mode_and_budget<R: SeekRead>(
+ reader: &mut IVFPQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ nprobe: usize,
+ reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_reader_filter_with_reuse_mode_and_budget(
+ reader,
+ queries,
+ nq,
+ k,
+ nprobe,
+ None,
+ reuse_mode,
+ reuse_max_bytes,
+ )
}
/// Big batch search with an optional row-id filter.
@@ -1732,6 +1809,28 @@ pub fn search_batch_reader_filter_with_reuse_mode<R:
SeekRead>(
nprobe: usize,
filter: Option<&dyn RowIdFilter>,
reuse_mode: IvfPqBatchTableReuseMode,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_reader_filter_with_reuse_mode_and_budget(
+ reader,
+ queries,
+ nq,
+ k,
+ nprobe,
+ filter,
+ reuse_mode,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ )
+}
+
+pub fn search_batch_reader_filter_with_reuse_mode_and_budget<R: SeekRead>(
+ reader: &mut IVFPQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ nprobe: usize,
+ filter: Option<&dyn RowIdFilter>,
+ reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
search_batch_reader_filter_with_reuse_mode_and_observer(
reader,
@@ -1741,7 +1840,10 @@ pub fn search_batch_reader_filter_with_reuse_mode<R:
SeekRead>(
nprobe,
filter,
reuse_mode,
+ reuse_max_bytes,
|_| {},
+ #[cfg(test)]
+ None,
)
}
@@ -1763,7 +1865,9 @@ fn search_batch_reader_filter_with_observer<R: SeekRead>(
nprobe,
filter,
IvfPqBatchTableReuseMode::Auto,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
&mut observe_ephemeral_precomputed_lists,
+ None,
)
}
@@ -1775,7 +1879,9 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
nprobe: usize,
filter: Option<&dyn RowIdFilter>,
reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
mut observe_ephemeral_precomputed_lists: impl FnMut(usize),
+ #[cfg(test)] distance_table_builds:
Option<&std::sync::atomic::AtomicUsize>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
reader.ensure_loaded()?;
let d = reader.d;
@@ -1856,10 +1962,17 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
}
unique_lists.sort_unstable_by_key(|&list_id| reader.list_offsets[list_id]);
+ let reuse_required_bytes = nq
+ .checked_mul(m)
+ .and_then(|values| values.checked_mul(ksub))
+ .and_then(|values| values.checked_mul(std::mem::size_of::<f32>()));
+ let reused_query_tables_fit_budget =
+ reuse_required_bytes.is_some_and(|bytes| bytes <= reuse_max_bytes);
let use_precomputed = reuse_mode != IvfPqBatchTableReuseMode::Off
&& metric == MetricType::L2
&& by_residual
- && !reader.precomputed_table.is_empty();
+ && !reader.precomputed_table.is_empty()
+ && reused_query_tables_fit_budget;
let allow_ephemeral_precomputed = reader.pq.nbits == 8
&& metric == MetricType::L2
&& by_residual
@@ -1883,6 +1996,36 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
} else {
Vec::new()
};
+ // 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
+ && !by_residual
+ && nprobe > 1
+ && match reuse_mode {
+ IvfPqBatchTableReuseMode::Off => false,
+ IvfPqBatchTableReuseMode::On => true,
+ IvfPqBatchTableReuseMode::Auto => {
+ let (active_query_count, probe_count) = all_probe_indices
+ .iter()
+ .map(|probes| {
+ probes
+ .iter()
+ .filter(|&&list_id| reader.list_counts[list_id] >
0)
+ .count()
+ })
+ .fold((0usize, 0usize), |(active, total), probes| {
+ (active + usize::from(probes > 0), total + probes)
+ });
+ nq >= MIN_EPHEMERAL_PRECOMPUTE_QUERIES
+ && should_use_ephemeral_precomputation(0,
active_query_count, probe_count)
+ }
+ }
+ && reused_query_tables_fit_budget;
+ let shared_sim_tables = if reuse_non_residual_tables {
+ (0..nq).map(|_| OnceLock::new()).collect::<Vec<_>>()
+ } else {
+ Vec::new()
+ };
let mut stable_pq_norms = None;
let mut heaps = (0..nq).map(|_| TopKHeap::new(k)).collect::<Vec<_>>();
@@ -1897,17 +2040,19 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
.position(|&list_id| list_id == first_list)
.map(|probe_rank| {
let query = &processed[query_index *
d..(query_index + 1) * d];
- let sim_table = reader_sim_table(
- reader,
- first_list,
- query,
- if use_precomputed {
- &all_ip_tables[query_index]
- } else {
- &[]
- },
- use_precomputed,
- );
+ let sim_table =
(!reuse_non_residual_tables).then(|| {
+ reader_sim_table(
+ reader,
+ first_list,
+ query,
+ if use_precomputed {
+ &all_ip_tables[query_index]
+ } else {
+ &[]
+ },
+ use_precomputed,
+ )
+ });
let dis0 = if use_precomputed {
all_coarse_dists[query_index][probe_rank]
} else {
@@ -1922,9 +2067,27 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
// The loop is sequential across queries. Reuse one chunk-sized
// distance buffer instead of retaining one per query.
let mut distances = Vec::new();
- reader.for_each_streamed_list_chunk(first_list, |ids, codes| {
+ reader.for_each_streamed_list_chunk(first_list, |pq, ids, codes| {
let positions = matching_rows(ids, filter);
+ if positions.as_ref().is_some_and(MatchingRows::is_empty) {
+ return;
+ }
for (query_index, dis0, sim_table) in &query_tables {
+ let sim_table = sim_table.as_deref().unwrap_or_else(|| {
+ shared_sim_tables[*query_index].get_or_init(|| {
+ let mut table = vec![0.0f32; m * ksub];
+ pq.compute_distance_table(
+ &processed[*query_index * d..(*query_index +
1) * d],
+ metric,
+ &mut table,
+ );
+ #[cfg(test)]
+ if let Some(builds) = distance_table_builds {
+ builds.fetch_add(1,
std::sync::atomic::Ordering::Relaxed);
+ }
+ table
+ })
+ });
scan_reader_codes(
sim_table,
codes,
@@ -1985,6 +2148,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
query_scratch_count,
m,
ksub,
+ reuse_max_bytes,
)
&& (reuse_mode == IvfPqBatchTableReuseMode::On
|| should_use_ephemeral_precomputation(
@@ -2033,6 +2197,9 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
&[]
},
use_precomputed,
+ shared_sim_table: reuse_non_residual_tables.then(||
&shared_sim_tables[qi]),
+ #[cfg(test)]
+ distance_table_builds,
d,
m,
ksub,
@@ -2122,6 +2289,28 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
}
}
+ if !by_residual &&
std::env::var_os("PAIMON_VINDEX_LOG_IVFPQ_BATCH_REUSE").is_some() {
+ use std::io::Write;
+
+ let tables_built = shared_sim_tables
+ .iter()
+ .filter(|table| table.get().is_some())
+ .count();
+ let _ = writeln!(
+ std::io::stderr().lock(),
+ "[paimon-vindex] ivfpq_batch_table_reuse
strategy=non_residual_query_table \
+ mode={reuse_mode:?} enabled={reuse_non_residual_tables} used={}
metric={} \
+ pq_bits={} nq={nq} nprobe={nprobe} unique_lists={} filtered={}
required_bytes={:?} \
+ budget_bytes={reuse_max_bytes} tables_built={tables_built}",
+ tables_built > 0,
+ metric.as_str(),
+ reader.pq.nbits,
+ unique_lists.len(),
+ filter.is_some(),
+ reuse_required_bytes,
+ );
+ }
+
Ok((result_ids, result_dists))
}
@@ -2153,9 +2342,31 @@ pub fn
search_batch_reader_roaring_filter_with_reuse_mode<R: SeekRead>(
nprobe: usize,
roaring_filter_bytes: &[u8],
reuse_mode: IvfPqBatchTableReuseMode,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_reader_roaring_filter_with_reuse_mode_and_budget(
+ reader,
+ queries,
+ nq,
+ k,
+ nprobe,
+ roaring_filter_bytes,
+ reuse_mode,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ )
+}
+
+pub fn search_batch_reader_roaring_filter_with_reuse_mode_and_budget<R:
SeekRead>(
+ reader: &mut IVFPQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ nprobe: usize,
+ roaring_filter_bytes: &[u8],
+ reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
let filter = decode_roaring_filter(roaring_filter_bytes)?;
- search_batch_reader_filter_with_reuse_mode(
+ search_batch_reader_filter_with_reuse_mode_and_budget(
reader,
queries,
nq,
@@ -2163,6 +2374,7 @@ pub fn
search_batch_reader_roaring_filter_with_reuse_mode<R: SeekRead>(
nprobe,
Some(&filter),
reuse_mode,
+ reuse_max_bytes,
)
}
@@ -2413,9 +2625,11 @@ mod tests {
nprobe,
filter,
reuse_mode,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
|count| {
precomputed_lists.fetch_add(count, Ordering::Relaxed);
},
+ None,
)
.unwrap();
@@ -3081,27 +3295,35 @@ mod tests {
#[test]
fn ephemeral_precomputation_respects_batch_memory_budget() {
- let max_values =
- crate::index_io_util::MAX_IVF_BATCH_READ_BYTES /
std::mem::size_of::<f64>();
+ let max_values = DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES /
std::mem::size_of::<f64>();
let max_list_values = max_values / 3;
assert!(ephemeral_precomputed_table_fits_budget(
1,
1,
1,
- max_list_values
+ max_list_values,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
));
assert!(!ephemeral_precomputed_table_fits_budget(
1,
1,
1,
- max_list_values + 1
+ max_list_values + 1,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ ));
+ assert!(!ephemeral_precomputed_table_fits_budget(
+ 0,
+ 1,
+ 1,
+ 1,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
));
- assert!(!ephemeral_precomputed_table_fits_budget(0, 1, 1, 1));
assert!(!ephemeral_precomputed_table_fits_budget(
usize::MAX,
usize::MAX,
usize::MAX,
- usize::MAX
+ usize::MAX,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
));
}
@@ -3571,6 +3793,154 @@ mod tests {
}
}
+ #[test]
+ fn
inner_product_batch_table_reuse_modes_preserve_results_and_control_table_builds()
{
+ use crate::io::{write_index, IVFPQIndexReader, PosWriter};
+
+ let d = 16;
+ let nlist = 4;
+ let m = 4;
+ let n = 600;
+ let nq = 8;
+ let k = 5;
+ let nprobe = nlist;
+ let data = generate_clustered_data(n, d, nlist, 43);
+ let ids = (0..n as i64).collect::<Vec<_>>();
+ let mut index = IVFPQIndex::new(d, nlist, m, MetricType::InnerProduct,
false);
+ index.train(&data, n);
+ index.add(&data, &ids, n);
+
+ let mut bytes = Vec::new();
+ write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+ let distance_table_builds = AtomicUsize::new(0);
+ let mut reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+
+ let (batch_ids, batch_distances) =
search_batch_reader_filter_with_reuse_mode_and_observer(
+ &mut reader,
+ &data[..nq * d],
+ nq,
+ k,
+ nprobe,
+ None,
+ IvfPqBatchTableReuseMode::On,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ |_| {},
+ Some(&distance_table_builds),
+ )
+ .unwrap();
+
+ assert_eq!(distance_table_builds.load(Ordering::Relaxed), nq);
+
+ let distance_table_builds = AtomicUsize::new(0);
+ let mut reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+ search_batch_reader_filter_with_reuse_mode_and_observer(
+ &mut reader,
+ &data[..nq * d],
+ nq,
+ k,
+ nprobe,
+ None,
+ IvfPqBatchTableReuseMode::On,
+ nq * m * 256 * std::mem::size_of::<f32>() - 1,
+ |_| {},
+ Some(&distance_table_builds),
+ )
+ .unwrap();
+ assert_eq!(
+ distance_table_builds.load(Ordering::Relaxed),
+ nq * nprobe,
+ "the direct path should be used when reused tables exceed the
configured budget"
+ );
+
+ let mut direct_reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+ let (direct_ids, direct_distances) =
search_batch_reader_with_reuse_mode(
+ &mut direct_reader,
+ &data[..nq * d],
+ nq,
+ k,
+ nprobe,
+ IvfPqBatchTableReuseMode::Off,
+ )
+ .unwrap();
+ assert_eq!(batch_ids, direct_ids);
+ assert_eq!(batch_distances, direct_distances);
+
+ for query_index in 0..nq {
+ let mut scalar_reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+ let query = &data[query_index * d..(query_index + 1) * d];
+ let (scalar_ids, scalar_distances) =
+ search_with_reader(&mut scalar_reader, query, k,
nprobe).unwrap();
+ let result = query_index * k..(query_index + 1) * k;
+ assert_eq!(&batch_ids[result.clone()], scalar_ids.as_slice());
+ assert_eq!(&batch_distances[result], scalar_distances.as_slice());
+ }
+
+ let large_nq = MIN_EPHEMERAL_PRECOMPUTE_QUERIES;
+ let distance_table_builds = AtomicUsize::new(0);
+ let mut reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+ search_batch_reader_filter_with_reuse_mode_and_observer(
+ &mut reader,
+ &data[..large_nq * d],
+ large_nq,
+ k,
+ nprobe,
+ None,
+ IvfPqBatchTableReuseMode::Auto,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ |_| {},
+ Some(&distance_table_builds),
+ )
+ .unwrap();
+ assert_eq!(
+ distance_table_builds.load(Ordering::Relaxed),
+ large_nq,
+ "Auto should reuse tables when the batch and probe work amortize
them"
+ );
+
+ let distance_table_builds = AtomicUsize::new(0);
+ let mut reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+ search_batch_reader_filter_with_reuse_mode_and_observer(
+ &mut reader,
+ &data[..nq * d],
+ nq,
+ k,
+ nprobe,
+ None,
+ IvfPqBatchTableReuseMode::Auto,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ |_| {},
+ Some(&distance_table_builds),
+ )
+ .unwrap();
+ assert_eq!(
+ distance_table_builds.load(Ordering::Relaxed),
+ nq * nprobe,
+ "Auto should keep the direct path for small batches"
+ );
+
+ let distance_table_builds = AtomicUsize::new(0);
+ let empty_filter = HashSet::new();
+ let mut reader = IVFPQIndexReader::open(Cursor::new(bytes)).unwrap();
+ search_batch_reader_filter_with_reuse_mode_and_observer(
+ &mut reader,
+ &data[..nq * d],
+ nq,
+ k,
+ nprobe,
+ Some(&empty_filter),
+ IvfPqBatchTableReuseMode::On,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ |_| {},
+ Some(&distance_table_builds),
+ )
+ .unwrap();
+ assert_eq!(
+ distance_table_builds.load(Ordering::Relaxed),
+ 0,
+ "empty filters should not build query distance tables"
+ );
+ }
+
#[test]
fn test_batch_reader_search_with_roaring_filter_bytes() {
use crate::io::{write_index, IVFPQIndexReader, PosWriter};
@@ -3985,7 +4355,7 @@ mod tests {
}
#[test]
- fn batch_table_reuse_off_ignores_resident_precomputed_tables() {
+ fn resident_precomputed_tables_respect_batch_reuse_mode_and_budget() {
use crate::io::{write_index, IVFPQIndexReader, PosWriter};
let d = 16;
@@ -4029,6 +4399,26 @@ mod tests {
.unwrap();
assert_eq!(actual, expected, "Off must ignore resident reuse tables");
+
+ let distance_table_builds = AtomicUsize::new(0);
+ let actual = search_batch_reader_filter_with_reuse_mode_and_observer(
+ &mut optimized_reader,
+ queries,
+ nq,
+ k,
+ nlist,
+ None,
+ IvfPqBatchTableReuseMode::On,
+ 1,
+ |_| {},
+ Some(&distance_table_builds),
+ )
+ .unwrap();
+ assert_eq!(
+ actual, expected,
+ "resident reuse tables must be ignored when query tables exceed
the budget"
+ );
+ assert_eq!(distance_table_builds.load(Ordering::Relaxed), nq * nlist);
}
#[test]
diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp
index d5a6a61..248c809 100644
--- a/cpp/test_vindex.cpp
+++ b/cpp/test_vindex.cpp
@@ -227,6 +227,8 @@ static void run_roundtrip(
auto batch_params = expected_index_type == PAIMON_VINDEX_INDEX_TYPE_DISKANN
? paimon::vindex::SearchParams::diskann(1, 100)
: paimon::vindex::SearchParams{1, 4};
+ batch_params.ivfpq_batch_table_reuse =
PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_OFF;
+ batch_params.ivfpq_batch_table_reuse_max_bytes = 1;
auto batch = reader.search_batch(queries.data(), 2, batch_params);
ASSERT_EQ(batch.ids.size(), 2);
assert_id_in_cluster(batch.ids[0], 0);
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 7a5368b..807c322 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -22,6 +22,7 @@ use paimon_vindex_core::index::{
IvfPqBatchTableReuseMode, SearchWidth, VectorIndexConfig,
VectorIndexMetadata,
VectorIndexReadPlan, VectorIndexReader, VectorIndexReaderOptions,
VectorIndexTrainer,
VectorIndexTraining, VectorIndexWriter, VectorSearchParams,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
};
use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities,
SeekWrite};
use std::cell::RefCell;
@@ -46,6 +47,11 @@ pub const PAIMON_VINDEX_SEARCH_WIDTH_AUTO: u32 = 0;
pub const PAIMON_VINDEX_SEARCH_WIDTH_IVF_NPROBE: u32 = 1;
pub const PAIMON_VINDEX_SEARCH_WIDTH_DISKANN_L_SEARCH: u32 = 2;
+pub const PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_OFF: u32 = 0;
+pub const PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_ON: u32 = 1;
+pub const PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO: u32 = 2;
+pub const PAIMON_VINDEX_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES: usize = 512
* 1024 * 1024;
+
thread_local! {
static LAST_ERROR: RefCell<Option<CString>> = const { RefCell::new(None) };
}
@@ -266,6 +272,16 @@ pub struct PaimonVindexSearchParams {
pub width: usize,
}
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct PaimonVindexSearchParamsV2 {
+ pub top_k: usize,
+ pub search_width: u32,
+ pub width: usize,
+ pub ivfpq_batch_table_reuse: u32,
+ pub ivfpq_batch_table_reuse_max_bytes: usize,
+}
+
#[repr(C)]
#[derive(Clone, Copy)]
pub struct PaimonVindexReaderOptions {
@@ -526,9 +542,31 @@ fn search_params_from_ffi(params:
PaimonVindexSearchParams) -> Result<VectorSear
search_width,
width: params.width,
ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto,
+ ivfpq_batch_table_reuse_max_bytes:
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
})
}
+fn search_params_v2_from_ffi(
+ params: PaimonVindexSearchParamsV2,
+) -> Result<VectorSearchParams, String> {
+ let mut result = search_params_from_ffi(PaimonVindexSearchParams {
+ top_k: params.top_k,
+ search_width: params.search_width,
+ width: params.width,
+ })?;
+ result.ivfpq_batch_table_reuse = match params.ivfpq_batch_table_reuse {
+ PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_OFF =>
IvfPqBatchTableReuseMode::Off,
+ PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_ON =>
IvfPqBatchTableReuseMode::On,
+ PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO =>
IvfPqBatchTableReuseMode::Auto,
+ value => return Err(format!("invalid IVF-PQ batch table reuse mode:
{value}")),
+ };
+ if params.ivfpq_batch_table_reuse_max_bytes == 0 {
+ return Err("IVF-PQ batch table reuse max bytes must be
positive".to_string());
+ }
+ result.ivfpq_batch_table_reuse_max_bytes =
params.ivfpq_batch_table_reuse_max_bytes;
+ Ok(result)
+}
+
// ======================== Trainer / Writer ========================
#[no_mangle]
@@ -940,6 +978,37 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch(
})
}
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_reader_search_batch_v2(
+ handle: *mut PaimonVindexReaderHandle,
+ queries: *const f32,
+ query_count: usize,
+ params: PaimonVindexSearchParamsV2,
+ out_ids: *mut i64,
+ out_distances: *mut f32,
+ result_len: usize,
+) -> c_int {
+ ffi_status(|| {
+ let handle = unsafe { reader_mut(handle) }?;
+ let query_len = checked_len(query_count, handle.inner.dimension(),
"queries")?;
+ let queries = unsafe { const_slice(queries, query_len, "queries") }?;
+ let params = search_params_v2_from_ffi(params)?;
+ let expected_len = checked_len(query_count, params.top_k, "batch
result")?;
+ let (ids, distances) = handle
+ .inner
+ .search_batch(queries, query_count, params)
+ .map_err(|e| format!("search_batch: {}", e))?;
+ copy_search_result(
+ &ids,
+ &distances,
+ out_ids,
+ out_distances,
+ result_len,
+ expected_len,
+ )
+ })
+}
+
#[no_mangle]
pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter(
handle: *mut PaimonVindexReaderHandle,
@@ -974,6 +1043,40 @@ pub unsafe extern "C" fn
paimon_vindex_reader_search_batch_with_roaring_filter(
})
}
+#[no_mangle]
+pub unsafe extern "C" fn
paimon_vindex_reader_search_batch_with_roaring_filter_v2(
+ handle: *mut PaimonVindexReaderHandle,
+ queries: *const f32,
+ query_count: usize,
+ params: PaimonVindexSearchParamsV2,
+ roaring_filter: *const u8,
+ roaring_filter_len: usize,
+ out_ids: *mut i64,
+ out_distances: *mut f32,
+ result_len: usize,
+) -> c_int {
+ ffi_status(|| {
+ let handle = unsafe { reader_mut(handle) }?;
+ let query_len = checked_len(query_count, handle.inner.dimension(),
"queries")?;
+ let queries = unsafe { const_slice(queries, query_len, "queries") }?;
+ let filter = unsafe { const_slice(roaring_filter, roaring_filter_len,
"roaring_filter") }?;
+ let params = search_params_v2_from_ffi(params)?;
+ let expected_len = checked_len(query_count, params.top_k, "batch
result")?;
+ let (ids, distances) = handle
+ .inner
+ .search_batch_with_roaring_filter(queries, query_count, params,
filter)
+ .map_err(|e| format!("search_batch_with_roaring_filter: {}", e))?;
+ copy_search_result(
+ &ids,
+ &distances,
+ out_ids,
+ out_distances,
+ result_len,
+ expected_len,
+ )
+ })
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -1085,6 +1188,10 @@ mod tests {
#[test]
fn ffi_search_parameters_preserve_diskann_width() {
+ assert_eq!(
+ PAIMON_VINDEX_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES
+ );
let params = search_params_from_ffi(PaimonVindexSearchParams {
top_k: 10,
search_width: PAIMON_VINDEX_SEARCH_WIDTH_DISKANN_L_SEARCH,
@@ -1098,5 +1205,35 @@ mod tests {
params.ivfpq_batch_table_reuse,
IvfPqBatchTableReuseMode::Auto
);
+ assert_eq!(
+ params.ivfpq_batch_table_reuse_max_bytes,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES
+ );
+ }
+
+ #[test]
+ fn ffi_v2_search_parameters_preserve_batch_reuse_options() {
+ let params = search_params_v2_from_ffi(PaimonVindexSearchParamsV2 {
+ top_k: 10,
+ search_width: PAIMON_VINDEX_SEARCH_WIDTH_IVF_NPROBE,
+ width: 16,
+ ivfpq_batch_table_reuse: PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_ON,
+ ivfpq_batch_table_reuse_max_bytes: 32 * 1024 * 1024,
+ })
+ .unwrap();
+
+ assert_eq!(params.ivfpq_batch_table_reuse,
IvfPqBatchTableReuseMode::On);
+ assert_eq!(params.ivfpq_batch_table_reuse_max_bytes, 32 * 1024 * 1024);
+
+ for (mode, max_bytes) in [(3, 1),
(PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO, 0)] {
+ assert!(search_params_v2_from_ffi(PaimonVindexSearchParamsV2 {
+ top_k: 10,
+ search_width: PAIMON_VINDEX_SEARCH_WIDTH_IVF_NPROBE,
+ width: 16,
+ ivfpq_batch_table_reuse: mode,
+ ivfpq_batch_table_reuse_max_bytes: max_bytes,
+ })
+ .is_err());
+ }
}
}
diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp
index 6fa01f1..0b3e8e6 100644
--- a/include/paimon_vindex.hpp
+++ b/include/paimon_vindex.hpp
@@ -228,6 +228,9 @@ struct SearchParams {
size_t top_k = 0;
uint32_t search_width = PAIMON_VINDEX_SEARCH_WIDTH_AUTO;
size_t width = 0;
+ uint32_t ivfpq_batch_table_reuse =
PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO;
+ size_t ivfpq_batch_table_reuse_max_bytes =
+ PAIMON_VINDEX_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES;
SearchParams(size_t top_k, size_t nprobe)
: top_k(top_k),
@@ -256,6 +259,16 @@ struct SearchParams {
return params;
}
+ PaimonVindexSearchParamsV2 to_ffi_v2() const {
+ PaimonVindexSearchParamsV2 params;
+ params.top_k = top_k;
+ params.search_width = search_width;
+ params.width = width;
+ params.ivfpq_batch_table_reuse = ivfpq_batch_table_reuse;
+ params.ivfpq_batch_table_reuse_max_bytes =
ivfpq_batch_table_reuse_max_bytes;
+ return params;
+ }
+
private:
SearchParams() = default;
};
@@ -603,11 +616,11 @@ public:
SearchResult result;
result.ids.resize(result_len);
result.distances.resize(result_len);
- check(paimon_vindex_reader_search_batch(
+ check(paimon_vindex_reader_search_batch_v2(
require_open(),
queries,
query_count,
- params.to_ffi(),
+ params.to_ffi_v2(),
result.ids.data(),
result.distances.data(),
result_len));
@@ -625,11 +638,11 @@ public:
SearchResult result;
result.ids.resize(result_len);
result.distances.resize(result_len);
- check(paimon_vindex_reader_search_batch_with_roaring_filter(
+ check(paimon_vindex_reader_search_batch_with_roaring_filter_v2(
require_open(),
queries,
query_count,
- params.to_ffi(),
+ params.to_ffi_v2(),
filter,
filter_len,
result.ids.data(),
diff --git
a/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java
b/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java
index ef5fdde..a491af5 100644
--- a/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java
+++ b/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java
@@ -19,6 +19,8 @@ package org.apache.paimon.index.vector;
public final class VectorSearchParams {
+ public static final long DEFAULT_IVF_PQ_BATCH_TABLE_REUSE_MAX_BYTES = 512L
* 1024 * 1024;
+
static final int SEARCH_WIDTH_AUTO = 0;
static final int SEARCH_WIDTH_IVF_NPROBE = 1;
static final int SEARCH_WIDTH_DISKANN_L_SEARCH = 2;
@@ -27,26 +29,37 @@ public final class VectorSearchParams {
private final int searchWidth;
private final int width;
private final int ivfPqBatchTableReuseMode;
+ private final long ivfPqBatchTableReuseMaxBytes;
public VectorSearchParams(int topK, int nprobe) {
this(
topK,
SEARCH_WIDTH_IVF_NPROBE,
nprobe,
- IvfPqBatchTableReuseMode.AUTO.code());
+ IvfPqBatchTableReuseMode.AUTO.code(),
+ DEFAULT_IVF_PQ_BATCH_TABLE_REUSE_MAX_BYTES);
}
private VectorSearchParams(
- int topK, int searchWidth, int width, int
ivfPqBatchTableReuseMode) {
+ int topK,
+ int searchWidth,
+ int width,
+ int ivfPqBatchTableReuseMode,
+ long ivfPqBatchTableReuseMaxBytes) {
this.topK = topK;
this.searchWidth = searchWidth;
this.width = width;
this.ivfPqBatchTableReuseMode = ivfPqBatchTableReuseMode;
+ this.ivfPqBatchTableReuseMaxBytes = ivfPqBatchTableReuseMaxBytes;
}
public static VectorSearchParams automatic(int topK) {
return new VectorSearchParams(
- topK, SEARCH_WIDTH_AUTO, 0,
IvfPqBatchTableReuseMode.AUTO.code());
+ topK,
+ SEARCH_WIDTH_AUTO,
+ 0,
+ IvfPqBatchTableReuseMode.AUTO.code(),
+ DEFAULT_IVF_PQ_BATCH_TABLE_REUSE_MAX_BYTES);
}
public static VectorSearchParams ivf(int topK, int nprobe) {
@@ -54,7 +67,8 @@ public final class VectorSearchParams {
topK,
SEARCH_WIDTH_IVF_NPROBE,
nprobe,
- IvfPqBatchTableReuseMode.AUTO.code());
+ IvfPqBatchTableReuseMode.AUTO.code(),
+ DEFAULT_IVF_PQ_BATCH_TABLE_REUSE_MAX_BYTES);
}
public static VectorSearchParams diskAnn(int topK, int lSearch) {
@@ -62,7 +76,8 @@ public final class VectorSearchParams {
topK,
SEARCH_WIDTH_DISKANN_L_SEARCH,
lSearch,
- IvfPqBatchTableReuseMode.AUTO.code());
+ IvfPqBatchTableReuseMode.AUTO.code(),
+ DEFAULT_IVF_PQ_BATCH_TABLE_REUSE_MAX_BYTES);
}
public int topK() {
@@ -85,19 +100,37 @@ public final class VectorSearchParams {
return ivfPqBatchTableReuseMode;
}
+ public long ivfPqBatchTableReuseMaxBytes() {
+ return ivfPqBatchTableReuseMaxBytes;
+ }
+
public VectorSearchParams
withIvfPqBatchTableReuse(IvfPqBatchTableReuseMode mode) {
if (mode == null) {
throw new IllegalArgumentException("IVF-PQ batch table reuse mode
is null");
}
- return new VectorSearchParams(topK, searchWidth, width, mode.code());
+ return new VectorSearchParams(
+ topK, searchWidth, width, mode.code(),
ivfPqBatchTableReuseMaxBytes);
}
public VectorSearchParams withIvfPqBatchTableReuse(String mode) {
return
withIvfPqBatchTableReuse(IvfPqBatchTableReuseMode.fromString(mode));
}
+ public VectorSearchParams withIvfPqBatchTableReuseMaxBytes(long maxBytes) {
+ if (maxBytes <= 0) {
+ throw new IllegalArgumentException(
+ "IVF-PQ batch table reuse max bytes must be positive");
+ }
+ return new VectorSearchParams(
+ topK, searchWidth, width, ivfPqBatchTableReuseMode, maxBytes);
+ }
+
public VectorSearchParams withLSearch(int lSearch) {
return new VectorSearchParams(
- topK, SEARCH_WIDTH_DISKANN_L_SEARCH, lSearch,
ivfPqBatchTableReuseMode);
+ topK,
+ SEARCH_WIDTH_DISKANN_L_SEARCH,
+ lSearch,
+ ivfPqBatchTableReuseMode,
+ ivfPqBatchTableReuseMaxBytes);
}
}
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 507d79c..fed0feb 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
@@ -74,6 +74,19 @@ public class VectorIndexJavaApiTest {
private static void testIvfPqBatchTableReuseMode() {
VectorSearchParams defaults = new VectorSearchParams(10, 4);
assertEquals(IvfPqBatchTableReuseMode.AUTO,
defaults.ivfPqBatchTableReuse());
+ assertEquals(512L * 1024 * 1024,
defaults.ivfPqBatchTableReuseMaxBytes());
+
+ VectorSearchParams customBudget =
+ defaults.withIvfPqBatchTableReuseMaxBytes(128L * 1024 * 1024);
+ assertEquals(128L * 1024 * 1024,
customBudget.ivfPqBatchTableReuseMaxBytes());
+ assertThrows(
+ IllegalArgumentException.class,
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ defaults.withIvfPqBatchTableReuseMaxBytes(0);
+ }
+ });
VectorSearchParams enabled =
defaults.withIvfPqBatchTableReuse(IvfPqBatchTableReuseMode.ON);
diff --git a/jni/src/lib.rs b/jni/src/lib.rs
index a189483..9fb9afa 100644
--- a/jni/src/lib.rs
+++ b/jni/src/lib.rs
@@ -415,6 +415,10 @@ fn search_params(env: &mut JNIEnv, params: JObject) ->
Result<VectorSearchParams
let width = call_int_method(env, ¶ms, "width")?;
let ivfpq_batch_table_reuse =
ivfpq_batch_table_reuse_mode(call_int_method(env, ¶ms,
"ivfPqBatchTableReuseMode")?)?;
+ let ivfpq_batch_table_reuse_max_bytes = positive_jlong_to_usize(
+ call_long_method(env, ¶ms, "ivfPqBatchTableReuseMaxBytes")?,
+ "IVF-PQ batch table reuse max bytes",
+ )?;
if top_k < 0 || width < 0 {
return Err(format!(
"invalid search parameters: topK={}, searchWidth={}, width={}",
@@ -435,6 +439,7 @@ fn search_params(env: &mut JNIEnv, params: JObject) ->
Result<VectorSearchParams
search_width,
width: width as usize,
ivfpq_batch_table_reuse,
+ ivfpq_batch_table_reuse_max_bytes,
})
}
@@ -453,6 +458,19 @@ fn call_int_method(env: &mut JNIEnv, object: &JObject,
name: &str) -> Result<jin
.map_err(|e| format!("VectorSearchParams.{}(): {}", name, e))
}
+fn call_long_method(env: &mut JNIEnv, object: &JObject, name: &str) ->
Result<jlong, String> {
+ env.call_method(object, name, "()J", &[])
+ .and_then(|value| value.j())
+ .map_err(|e| format!("VectorSearchParams.{}(): {}", name, e))
+}
+
+fn positive_jlong_to_usize(value: jlong, name: &str) -> Result<usize, String> {
+ if value <= 0 {
+ return Err(format!("{name} must be positive"));
+ }
+ usize::try_from(value).map_err(|_| format!("{name} exceeds usize"))
+}
+
// --- Unified Trainer / Writer API ---
#[no_mangle]
@@ -1056,4 +1074,14 @@ mod tests {
);
assert!(ivfpq_batch_table_reuse_mode(3).is_err());
}
+
+ #[test]
+ fn ivfpq_batch_table_reuse_budget_must_be_positive() {
+ assert_eq!(
+ positive_jlong_to_usize(64 * 1024 * 1024, "reuse max
bytes").unwrap(),
+ 64 * 1024 * 1024
+ );
+ assert!(positive_jlong_to_usize(0, "reuse max bytes").is_err());
+ assert!(positive_jlong_to_usize(-1, "reuse max bytes").is_err());
+ }
}
diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py
index 99037b9..bba1ef0 100644
--- a/python/paimon_vindex/__init__.py
+++ b/python/paimon_vindex/__init__.py
@@ -29,6 +29,7 @@ from ._ffi import lib
_SIZE_T_MAX = ctypes.c_size_t(-1).value
_UINT64_MAX = ctypes.c_uint64(-1).value
+_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES = 512 * 1024 * 1024
def _size_t(value, name: str, *, allow_zero: bool) -> int:
@@ -121,6 +122,12 @@ class SearchWidth(IntEnum):
DISKANN_L_SEARCH = 2
+class IvfPqBatchTableReuseMode(IntEnum):
+ OFF = 0
+ ON = 1
+ AUTO = 2
+
+
@dataclass(frozen=True)
class VectorIndexMetadata:
index_type: str
@@ -154,12 +161,22 @@ class SearchParams:
top_k: int
search_width: SearchWidth = SearchWidth.AUTO
width: int = 0
+ ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode = (
+ IvfPqBatchTableReuseMode.AUTO
+ )
+ ivfpq_batch_table_reuse_max_bytes: int = (
+ _DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES
+ )
def __post_init__(self):
try:
search_width = SearchWidth(self.search_width)
except (TypeError, ValueError) as exc:
raise ValueError("search_width is invalid") from exc
+ try:
+ reuse_mode = IvfPqBatchTableReuseMode(self.ivfpq_batch_table_reuse)
+ except (TypeError, ValueError) as exc:
+ raise ValueError("IVF-PQ batch table reuse mode is invalid") from
exc
top_k = _size_t(self.top_k, "top_k", allow_zero=False)
if search_width == SearchWidth.AUTO:
width = _size_t(self.width, "automatic search width",
allow_zero=True)
@@ -172,20 +189,37 @@ class SearchParams:
else "l_search"
)
width = _size_t(self.width, name, allow_zero=False)
+ reuse_max_bytes = _size_t(
+ self.ivfpq_batch_table_reuse_max_bytes,
+ "IVF-PQ batch table reuse max bytes",
+ allow_zero=False,
+ )
object.__setattr__(self, "top_k", top_k)
object.__setattr__(self, "search_width", search_width)
object.__setattr__(self, "width", width)
+ object.__setattr__(self, "ivfpq_batch_table_reuse", reuse_mode)
+ object.__setattr__(
+ self, "ivfpq_batch_table_reuse_max_bytes", reuse_max_bytes
+ )
@classmethod
def automatic(cls, top_k: int):
return cls(top_k=top_k)
@classmethod
- def ivf(cls, top_k: int, nprobe: int):
+ def ivf(
+ cls,
+ top_k: int,
+ nprobe: int,
+ ivfpq_batch_table_reuse=IvfPqBatchTableReuseMode.AUTO,
+
ivfpq_batch_table_reuse_max_bytes=_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ ):
return cls(
top_k=top_k,
search_width=SearchWidth.IVF_NPROBE,
width=nprobe,
+ ivfpq_batch_table_reuse=ivfpq_batch_table_reuse,
+
ivfpq_batch_table_reuse_max_bytes=ivfpq_batch_table_reuse_max_bytes,
)
@classmethod
@@ -203,6 +237,15 @@ class SearchParams:
self.width,
)
+ def to_ffi_v2(self):
+ return _ffi.PaimonVindexSearchParamsV2(
+ self.top_k,
+ int(self.search_width),
+ self.width,
+ int(self.ivfpq_batch_table_reuse),
+ self.ivfpq_batch_table_reuse_max_bytes,
+ )
+
def _check_error(message="operation failed"):
err = lib.paimon_vindex_last_error()
@@ -774,7 +817,7 @@ class VectorIndexReader:
f"queries length {queries.size} does not match nq * dimension "
f"{queries.shape[0] * self._metadata.dimension}"
)
- ffi_params = params.to_ffi()
+ ffi_params = params.to_ffi_v2()
result_len = queries.shape[0] * params.top_k
ids = np.empty((queries.shape[0], params.top_k), dtype=np.int64)
distances = np.empty((queries.shape[0], params.top_k),
dtype=np.float32)
@@ -782,7 +825,7 @@ class VectorIndexReader:
with self._native_handle_lock:
self._require_open()
if filter_bytes is None:
- rc = lib.paimon_vindex_reader_search_batch(
+ rc = lib.paimon_vindex_reader_search_batch_v2(
self._handle,
queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
queries.shape[0],
@@ -793,7 +836,7 @@ class VectorIndexReader:
)
else:
filter_buf, filter_len, _ = self._filter_args(filter_bytes)
- rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter(
+ rc =
lib.paimon_vindex_reader_search_batch_with_roaring_filter_v2(
self._handle,
queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
queries.shape[0],
@@ -832,6 +875,7 @@ class VectorIndexReader:
__all__ = [
+ "IvfPqBatchTableReuseMode",
"SearchParams",
"VectorIndexMetadata",
"VectorIndexReadPlan",
diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py
index f86936f..02cd900 100644
--- a/python/paimon_vindex/_ffi.py
+++ b/python/paimon_vindex/_ffi.py
@@ -146,6 +146,16 @@ class PaimonVindexSearchParams(Structure):
]
+class PaimonVindexSearchParamsV2(Structure):
+ _fields_ = [
+ ("top_k", c_size_t),
+ ("search_width", c_uint32),
+ ("width", c_size_t),
+ ("ivfpq_batch_table_reuse", c_uint32),
+ ("ivfpq_batch_table_reuse_max_bytes", c_size_t),
+ ]
+
+
class PaimonVindexReaderOptions(Structure):
_fields_ = [
("memory_budget_bytes", c_size_t),
@@ -292,6 +302,17 @@ lib.paimon_vindex_reader_search_batch.argtypes = [
]
lib.paimon_vindex_reader_search_batch.restype = c_int
+lib.paimon_vindex_reader_search_batch_v2.argtypes = [
+ c_void_p,
+ POINTER(c_float),
+ c_size_t,
+ PaimonVindexSearchParamsV2,
+ POINTER(c_int64),
+ POINTER(c_float),
+ c_size_t,
+]
+lib.paimon_vindex_reader_search_batch_v2.restype = c_int
+
lib.paimon_vindex_reader_search_batch_with_roaring_filter.argtypes = [
c_void_p,
POINTER(c_float),
@@ -304,3 +325,16 @@
lib.paimon_vindex_reader_search_batch_with_roaring_filter.argtypes = [
c_size_t,
]
lib.paimon_vindex_reader_search_batch_with_roaring_filter.restype = c_int
+
+lib.paimon_vindex_reader_search_batch_with_roaring_filter_v2.argtypes = [
+ c_void_p,
+ POINTER(c_float),
+ c_size_t,
+ PaimonVindexSearchParamsV2,
+ POINTER(c_uint8),
+ c_size_t,
+ POINTER(c_int64),
+ POINTER(c_float),
+ c_size_t,
+]
+lib.paimon_vindex_reader_search_batch_with_roaring_filter_v2.restype = c_int
diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py
index 8c638f4..950f714 100644
--- a/python/tests/test_vindex.py
+++ b/python/tests/test_vindex.py
@@ -23,6 +23,7 @@ import numpy as np
import pytest
from paimon_vindex import (
+ IvfPqBatchTableReuseMode,
SearchParams,
VectorIndexReader,
VectorIndexTrainer,
@@ -70,6 +71,19 @@ def
test_python_search_parameters_remain_algorithm_specific():
automatic = SearchParams.automatic(top_k=10).to_ffi()
assert automatic.search_width == 0
assert automatic.width == 0
+ assert (
+ SearchParams.ivf(10, 16).ivfpq_batch_table_reuse_max_bytes
+ == 512 * 1024 * 1024
+ )
+
+ batch = SearchParams.ivf(
+ top_k=10,
+ nprobe=16,
+ ivfpq_batch_table_reuse=IvfPqBatchTableReuseMode.ON,
+ ivfpq_batch_table_reuse_max_bytes=32 * 1024 * 1024,
+ ).to_ffi_v2()
+ assert batch.ivfpq_batch_table_reuse == 1
+ assert batch.ivfpq_batch_table_reuse_max_bytes == 32 * 1024 * 1024
@pytest.mark.parametrize(
@@ -81,6 +95,14 @@ def
test_python_search_parameters_remain_algorithm_specific():
lambda: SearchParams.ivf(
top_k=5, nprobe=ctypes.c_size_t(-1).value + 1
),
+ lambda: SearchParams.ivf(
+ top_k=5, nprobe=2, ivfpq_batch_table_reuse=3
+ ),
+ lambda: SearchParams.ivf(
+ top_k=5,
+ nprobe=2,
+ ivfpq_batch_table_reuse_max_bytes=0,
+ ),
],
)
def
test_python_search_parameters_reject_values_that_ctypes_would_wrap(factory):
