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 554b121 Expose optimize for search API (#40)
554b121 is described below
commit 554b12143f2e8de1c3975eaabd778dcc86b81e98
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Jun 12 10:19:14 2026 +0800
Expose optimize for search API (#40)
---
README.md | 16 +++
c/test_vindex.c | 4 +
core/src/index.rs | 79 ++++++++++++++
core/src/io.rs | 9 ++
core/src/ivfhnswsq_io.rs | 87 ++++++++++++++-
core/src/sq.rs | 119 +++++++++++++++++++++
cpp/test_vindex.cpp | 2 +
ffi/src/lib.rs | 13 +++
include/paimon_vindex.hpp | 4 +
.../paimon/index/vector/VectorIndexNative.java | 2 +
.../paimon/index/vector/VectorIndexReader.java | 11 ++
.../index/vector/VectorIndexJavaApiTest.java | 7 ++
jni/src/lib.rs | 17 +++
python/paimon_vindex/__init__.py | 6 ++
python/paimon_vindex/_ffi.py | 3 +
python/tests/test_vindex.py | 6 ++
16 files changed, 380 insertions(+), 5 deletions(-)
diff --git a/README.md b/README.md
index 393a61e..3ea1782 100644
--- a/README.md
+++ b/README.md
@@ -68,6 +68,17 @@ types:
- `ef_search`: optional HNSW search breadth for `IVF_HNSW_FLAT` and
`IVF_HNSW_SQ`. A value of `0` uses the default.
+Readers also expose an optional search warm-up API. Call
+`optimize_for_search` in Rust, C++, and Python,
+`paimon_vindex_reader_optimize_for_search` in the C ABI, or
+`optimizeForSearch` in Java after opening a reader and before repeated
+searches. The call builds in-memory search caches and does not change the
+serialized index format or search results. Currently `IVF_PQ` builds residual
+L2 precomputed tables for repeated PQ searches, `IVF_HNSW_SQ` builds SQ decode
+LUTs for filtered-search SQ scan and fallback paths, and other index types
+preload metadata. The `IVF_HNSW_SQ` LUTs are not expected to speed up the
+normal unfiltered HNSW graph-search path.
+
### Rust
```rust
@@ -97,6 +108,7 @@ writer.write(&mut out)?;
let file = File::open("vectors.pvindex")?;
let mut reader = VectorIndexReader::open(file)?;
+reader.optimize_for_search()?;
let params = VectorSearchParams::with_ef_search(10, 16, 80);
let (ids, distances) = reader.search(&query, params)?;
```
@@ -148,6 +160,7 @@ paimon_vindex_writer_free(writer);
PaimonVindexReaderHandle *reader = paimon_vindex_reader_open(input_file);
PaimonVindexMetadata metadata;
paimon_vindex_reader_metadata(reader, &metadata);
+paimon_vindex_reader_optimize_for_search(reader);
int64_t ids[10];
float distances[10];
@@ -183,6 +196,7 @@ writer.write_index(output_file);
paimon::vindex::Reader reader(input_file);
auto metadata = reader.metadata();
+reader.optimize_for_search();
auto result = reader.search(query.data(), 10, 16, 80);
```
@@ -215,6 +229,7 @@ try (VectorIndexWriter writer = new
VectorIndexWriter(options)) {
try (VectorIndexReader reader = new VectorIndexReader(vectorIndexInput)) {
VectorIndexMetadata metadata = reader.metadata();
+ reader.optimizeForSearch();
VectorSearchResult result = reader.search(query, 10, 16, 80);
}
```
@@ -252,6 +267,7 @@ writer.add_vectors(row_ids, vectors)
writer.write(output)
reader = VectorIndexReader(VectorIndexInput(index_bytes))
+reader.optimize_for_search()
ids, distances = reader.search(query, top_k=10, nprobe=16, ef_search=80)
```
diff --git a/c/test_vindex.c b/c/test_vindex.c
index 816dc80..916889e 100644
--- a/c/test_vindex.c
+++ b/c/test_vindex.c
@@ -166,6 +166,10 @@ static void test_basic_roundtrip(void) {
ASSERT_EQ_I64(metadata.nlist, 2);
ASSERT_EQ_I64(metadata.total_vectors, 4);
+ if (paimon_vindex_reader_optimize_for_search(reader) != 0) {
+ fail_ffi("reader optimize_for_search failed");
+ }
+
const float query[] = {0.0f, 0.0f};
int64_t result_ids[2] = {0};
float result_distances[2] = {0};
diff --git a/core/src/index.rs b/core/src/index.rs
index e66796b..c0b890c 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -526,6 +526,17 @@ impl<R: SeekRead> VectorIndexReader<R> {
self.metadata().total_vectors
}
+ pub fn optimize_for_search(&mut self) -> io::Result<()> {
+ match self {
+ Self::IvfFlat(reader) => reader.ensure_loaded(),
+ Self::IvfPq(reader) => reader.optimize_for_search(),
+ Self::IvfHnswFlat(reader) => reader.ensure_loaded(),
+ // IVF_HNSW_SQ warms SQ scan/fallback structures used by filtered
+ // searches; normal unfiltered search primarily uses the HNSW
graph.
+ Self::IvfHnswSq(reader) => reader.optimize_for_search(),
+ }
+ }
+
pub fn search(
&mut self,
query: &[f32],
@@ -824,6 +835,22 @@ mod tests {
assert_eq!(result_ids[0], 0);
}
+ fn build_reader(config: VectorIndexConfig) ->
(VectorIndexReader<Cursor<Vec<u8>>>, Vec<f32>) {
+ let d = config.dimension();
+ let nlist = config.nlist();
+ let n = 512;
+ let data = generate_clustered_data(n, d, nlist);
+ let ids = (0..n as i64).collect::<Vec<_>>();
+
+ let mut writer = VectorIndexWriter::new(config).unwrap();
+ writer.train(&data, n).unwrap();
+ writer.add_vectors(&ids, &data, n).unwrap();
+
+ let mut buf = Vec::new();
+ writer.write(&mut PosWriter::new(&mut buf)).unwrap();
+ (VectorIndexReader::open(Cursor::new(buf)).unwrap(), data)
+ }
+
fn build_ivfflat_reader() -> VectorIndexReader<Cursor<Vec<u8>>> {
let mut writer = VectorIndexWriter::new(VectorIndexConfig::IvfFlat {
dimension: 1,
@@ -878,6 +905,58 @@ mod tests {
});
}
+ #[test]
+ fn optimize_for_search_preserves_results() {
+ for config in [
+ VectorIndexConfig::IvfFlat {
+ dimension: 8,
+ nlist: 4,
+ metric: MetricType::L2,
+ },
+ VectorIndexConfig::IvfPq {
+ dimension: 16,
+ nlist: 4,
+ m: 4,
+ metric: MetricType::L2,
+ use_opq: false,
+ },
+ VectorIndexConfig::IvfHnswFlat {
+ dimension: 8,
+ nlist: 4,
+ metric: MetricType::L2,
+ hnsw: HnswBuildParams::default(),
+ },
+ VectorIndexConfig::IvfHnswSq {
+ dimension: 8,
+ nlist: 4,
+ metric: MetricType::L2,
+ hnsw: HnswBuildParams::default(),
+ },
+ ] {
+ let d = config.dimension();
+ let nlist = config.nlist();
+ let params = VectorSearchParams::with_ef_search(5, nlist, 32);
+ let (mut baseline, data) = build_reader(config.clone());
+ let query = data[0..d].to_vec();
+ let expected = baseline.search(&query, params).unwrap();
+
+ let (mut optimized, _) = build_reader(config);
+ optimized.optimize_for_search().unwrap();
+ let actual = optimized.search(&query, params).unwrap();
+
+ assert_eq!(actual.0, expected.0);
+ assert_eq!(actual.1.len(), expected.1.len());
+ for (actual, expected) in actual.1.iter().zip(expected.1.iter()) {
+ assert!(
+ (actual - expected).abs() < 1e-4,
+ "optimized distance {} should match baseline {}",
+ actual,
+ expected
+ );
+ }
+ }
+ }
+
#[test]
fn unified_reader_rejects_unknown_magic() {
let err = match VectorIndexReader::open(Cursor::new(vec![0xFF; 8])) {
diff --git a/core/src/io.rs b/core/src/io.rs
index 7570c82..60772de 100644
--- a/core/src/io.rs
+++ b/core/src/io.rs
@@ -648,6 +648,15 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
Ok(())
}
+ pub fn optimize_for_search(&mut self) -> io::Result<()> {
+ self.ensure_loaded()?;
+ if self.metric == MetricType::L2 && self.by_residual &&
self.precomputed_table.is_empty() {
+ self.precomputed_table =
+ compute_precomputed_table(&self.quantizer_centroids, &self.pq,
self.nlist, self.d);
+ }
+ Ok(())
+ }
+
/// Read an inverted list's IDs and PQ codes.
/// Calls ensure_loaded() if not yet loaded.
pub fn read_inverted_list(&mut self, list_id: usize) ->
io::Result<(Vec<i64>, Vec<u8>)> {
diff --git a/core/src/ivfhnswsq_io.rs b/core/src/ivfhnswsq_io.rs
index d5ea2dc..af159bc 100644
--- a/core/src/ivfhnswsq_io.rs
+++ b/core/src/ivfhnswsq_io.rs
@@ -29,9 +29,10 @@ use crate::io::{PreadCursor, ReadRequest, SeekRead,
SeekWrite};
use crate::ivfhnswsq::IVFHNSWSQIndex;
use crate::ivfpq::RowIdFilter;
use crate::kmeans;
-use crate::sq::ScalarQuantizer;
+use crate::sq::{ScalarQuantizer, ScalarQuantizerDecodeLut};
use crate::topk::TopKHeap;
use std::io;
+use std::sync::Arc;
pub const IVF_HNSW_SQ_MAGIC: u32 = 0x49485351; // "IHSQ"
pub const IVF_HNSW_SQ_VERSION: u32 = 1;
@@ -163,6 +164,7 @@ pub struct IVFHNSWSQIndexReader<R: SeekRead> {
pub list_counts: Vec<i32>,
pub list_graph_bytes_lens: Vec<i32>,
pub list_payload_bytes_lens: Vec<i64>,
+ sq_decode_luts: Vec<Arc<ScalarQuantizerDecodeLut>>,
loaded: bool,
}
@@ -259,6 +261,7 @@ impl<R: SeekRead> IVFHNSWSQIndexReader<R> {
list_counts: Vec::new(),
list_graph_bytes_lens: Vec::new(),
list_payload_bytes_lens: Vec::new(),
+ sq_decode_luts: Vec::new(),
loaded: false,
})
}
@@ -315,6 +318,21 @@ impl<R: SeekRead> IVFHNSWSQIndexReader<R> {
Ok(())
}
+ pub fn optimize_for_search(&mut self) -> io::Result<()> {
+ self.ensure_loaded()?;
+ if self.sq_decode_luts.len() != self.nlist {
+ // These LUTs only help when search falls back to scanning SQ
codes,
+ // for example filtered searches with a small candidate set. The
+ // normal unfiltered path searches decoded vectors in the HNSW
graph.
+ self.sq_decode_luts = self
+ .list_sqs
+ .iter()
+ .map(|sq| Arc::new(sq.build_decode_lut()))
+ .collect();
+ }
+ Ok(())
+ }
+
pub fn read_inverted_list(
&mut self,
list_id: usize,
@@ -588,6 +606,7 @@ impl<R: SeekRead> IVFHNSWSQIndexReader<R> {
graph,
centroid: Some(centroid),
sq: self.list_sq(meta.list_id).clone(),
+ sq_decode_lut:
self.list_sq_decode_lut(meta.list_id).map(Arc::clone),
})
}
@@ -599,6 +618,10 @@ impl<R: SeekRead> IVFHNSWSQIndexReader<R> {
self.list_sqs.get(list_id).unwrap_or(&self.sq)
}
+ fn list_sq_decode_lut(&self, list_id: usize) ->
Option<&Arc<ScalarQuantizerDecodeLut>> {
+ self.sq_decode_luts.get(list_id)
+ }
+
pub fn search(
&mut self,
query: &[f32],
@@ -643,6 +666,7 @@ impl<R: SeekRead> IVFHNSWSQIndexReader<R> {
&list.codes,
list.centroid.as_deref(),
&list.sq,
+ list.sq_decode_lut.as_deref(),
self.metric,
filter,
heap,
@@ -736,6 +760,7 @@ pub fn search_batch_ivfhnswsq_reader_filter<R: SeekRead>(
graph: list.graph,
centroid: list.centroid,
sq: list.sq,
+ sq_decode_lut: list.sq_decode_lut,
});
}
@@ -752,6 +777,7 @@ pub fn search_batch_ivfhnswsq_reader_filter<R: SeekRead>(
&list.codes,
list.centroid.as_deref(),
&list.sq,
+ list.sq_decode_lut.as_deref(),
reader.metric,
filter,
&mut heaps[qi],
@@ -783,6 +809,7 @@ pub fn search_batch_ivfhnswsq_reader_filter<R: SeekRead>(
&list.codes,
list.centroid.as_deref(),
&list.sq,
+ list.sq_decode_lut.as_deref(),
reader.metric,
filter,
&mut heaps[qi],
@@ -824,6 +851,7 @@ struct GraphList {
graph: HnswGraph,
centroid: Option<Vec<f32>>,
sq: ScalarQuantizer,
+ sq_decode_lut: Option<Arc<ScalarQuantizerDecodeLut>>,
}
#[derive(Clone, Copy)]
@@ -873,6 +901,7 @@ struct LoadedBatchList {
graph: HnswGraph,
centroid: Option<Vec<f32>>,
sq: ScalarQuantizer,
+ sq_decode_lut: Option<Arc<ScalarQuantizerDecodeLut>>,
}
fn scan_sq_list(
@@ -881,6 +910,7 @@ fn scan_sq_list(
codes: &[u8],
centroid: Option<&[f32]>,
sq: &ScalarQuantizer,
+ sq_decode_lut: Option<&ScalarQuantizerDecodeLut>,
metric: MetricType,
filter: Option<&dyn RowIdFilter>,
heap: &mut TopKHeap,
@@ -892,10 +922,16 @@ fn scan_sq_list(
continue;
}
let code = &codes[local_id * code_size..(local_id + 1) * code_size];
- let dist = if let Some(centroid) = centroid {
- sq.distance_to_code_with_offset_with_context(query, code,
centroid, context)
- } else {
- sq.distance_to_code_with_context(query, code, context)
+ let dist = match (centroid, sq_decode_lut) {
+ (Some(centroid), Some(lut)) => sq
+ .distance_to_code_with_lut_offset_with_context(query, code,
centroid, lut, context),
+ (Some(centroid), None) => {
+ sq.distance_to_code_with_offset_with_context(query, code,
centroid, context)
+ }
+ (None, Some(lut)) => {
+ sq.distance_to_code_with_lut_with_context(query, code, lut,
context)
+ }
+ (None, None) => sq.distance_to_code_with_context(query, code,
context),
};
heap.push(dist, row_id);
}
@@ -1384,6 +1420,47 @@ mod tests {
assert_eq!(labels, vec![12, -1]);
}
+ #[test]
+ fn test_ivfhnswsq_reader_optimized_filter_search_matches_unoptimized() {
+ let d = 4;
+ let nlist = 4;
+ let n = 128;
+ let data: Vec<f32> = (0..n)
+ .flat_map(|i| {
+ let cluster = (i % nlist) as f32 * 100.0;
+ [cluster + i as f32 * 0.01, 1.0, 2.0, 3.0]
+ })
+ .collect();
+ let ids: Vec<i64> = (0..n as i64).collect();
+
+ let mut index = IVFHNSWSQIndex::new(d, nlist, MetricType::L2,
HnswBuildParams::default());
+ index.train(&data, n);
+ index.add(&data, &ids, n);
+ index.build_graphs().unwrap();
+
+ let mut buf = Vec::new();
+ write_ivfhnswsq_index(&index, &mut PosWriter::new(&mut buf)).unwrap();
+
+ let mut filter = RoaringTreemap::new();
+ filter.insert(0);
+ filter.insert(64);
+ let mut filter_bytes = Vec::new();
+ filter.serialize_into(&mut filter_bytes).unwrap();
+
+ let mut baseline =
IVFHNSWSQIndexReader::open(Cursor::new(buf.clone())).unwrap();
+ let expected = baseline
+ .search_with_roaring_filter(&data[0..d], 3, nlist, 64,
&filter_bytes)
+ .unwrap();
+
+ let mut optimized =
IVFHNSWSQIndexReader::open(Cursor::new(buf)).unwrap();
+ optimized.optimize_for_search().unwrap();
+ let actual = optimized
+ .search_with_roaring_filter(&data[0..d], 3, nlist, 64,
&filter_bytes)
+ .unwrap();
+
+ assert_eq!(actual, expected);
+ }
+
#[test]
fn test_ivfhnswsq_reader_search_coalesces_contiguous_list_reads() {
let d = 4;
diff --git a/core/src/sq.rs b/core/src/sq.rs
index 812f94f..fe29509 100644
--- a/core/src/sq.rs
+++ b/core/src/sq.rs
@@ -26,6 +26,12 @@ pub struct ScalarQuantizer {
pub maxs: Vec<f32>,
}
+#[derive(Debug, Clone, PartialEq)]
+pub struct ScalarQuantizerDecodeLut {
+ d: usize,
+ values: Vec<f32>,
+}
+
impl ScalarQuantizer {
pub fn new(d: usize) -> Self {
Self {
@@ -162,6 +168,42 @@ impl ScalarQuantizer {
self.distance_to_code_impl(query, code, offset, true, context)
}
+ pub fn distance_to_code_with_lut_with_context(
+ &self,
+ query: &[f32],
+ code: &[u8],
+ lut: &ScalarQuantizerDecodeLut,
+ context: DistanceContext,
+ ) -> f32 {
+ self.distance_to_code_lut_impl(query, code, &[], false, lut, context)
+ }
+
+ pub fn distance_to_code_with_lut_offset_with_context(
+ &self,
+ query: &[f32],
+ code: &[u8],
+ offset: &[f32],
+ lut: &ScalarQuantizerDecodeLut,
+ context: DistanceContext,
+ ) -> f32 {
+ debug_assert!(query.len() >= self.d);
+ debug_assert!(code.len() >= self.d);
+ debug_assert!(offset.len() >= self.d);
+
+ self.distance_to_code_lut_impl(query, code, offset, true, lut, context)
+ }
+
+ pub fn build_decode_lut(&self) -> ScalarQuantizerDecodeLut {
+ let mut values = vec![0.0f32; self.d * 256];
+ for dim in 0..self.d {
+ let base = dim * 256;
+ for code in 0..256 {
+ values[base + code] = self.decode_value(code as u8, dim);
+ }
+ }
+ ScalarQuantizerDecodeLut { d: self.d, values }
+ }
+
fn distance_to_code_impl(
&self,
query: &[f32],
@@ -208,6 +250,55 @@ impl ScalarQuantizer {
}
}
+ fn distance_to_code_lut_impl(
+ &self,
+ query: &[f32],
+ code: &[u8],
+ offset: &[f32],
+ use_offset: bool,
+ lut: &ScalarQuantizerDecodeLut,
+ context: DistanceContext,
+ ) -> f32 {
+ debug_assert!(query.len() >= self.d);
+ debug_assert!(code.len() >= self.d);
+ debug_assert_eq!(lut.d, self.d);
+
+ match context.metric {
+ MetricType::L2 => {
+ let mut sum = 0.0f32;
+ for i in 0..self.d {
+ let diff = query[i]
+ - decoded_lut_value_with_offset(lut, code[i], i,
offset, use_offset);
+ sum += diff * diff;
+ }
+ sum
+ }
+ MetricType::InnerProduct => {
+ let mut dot = 0.0f32;
+ for i in 0..self.d {
+ dot += query[i]
+ * decoded_lut_value_with_offset(lut, code[i], i,
offset, use_offset);
+ }
+ -dot
+ }
+ MetricType::Cosine => {
+ let mut dot = 0.0f32;
+ let mut vector_norm = 0.0f32;
+ for i in 0..self.d {
+ let value = decoded_lut_value_with_offset(lut, code[i], i,
offset, use_offset);
+ dot += query[i] * value;
+ vector_norm += value * value;
+ }
+ let denom = context.query_norm * vector_norm.sqrt();
+ if denom > 0.0 {
+ 1.0 - dot / denom
+ } else {
+ 1.0
+ }
+ }
+ }
+ }
+
fn decode_value_with_offset(
&self,
code: u8,
@@ -253,6 +344,34 @@ impl ScalarQuantizer {
}
}
+impl ScalarQuantizerDecodeLut {
+ #[inline]
+ pub fn decode_value(&self, code: u8, dim: usize) -> f32 {
+ debug_assert!(dim < self.d);
+ self.values[dim * 256 + code as usize]
+ }
+
+ pub fn dimension(&self) -> usize {
+ self.d
+ }
+}
+
+#[inline]
+fn decoded_lut_value_with_offset(
+ lut: &ScalarQuantizerDecodeLut,
+ code: u8,
+ dim: usize,
+ offset: &[f32],
+ use_offset: bool,
+) -> f32 {
+ let value = lut.decode_value(code, dim);
+ if use_offset {
+ value + offset[dim]
+ } else {
+ value
+ }
+}
+
#[derive(Debug, Clone, Copy)]
pub struct DistanceContext {
metric: MetricType,
diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp
index 1b95840..c1eeb04 100644
--- a/cpp/test_vindex.cpp
+++ b/cpp/test_vindex.cpp
@@ -99,6 +99,8 @@ static void test_basic_roundtrip() {
ASSERT_EQ(metadata.metric, PAIMON_VINDEX_METRIC_L2);
ASSERT_EQ(metadata.total_vectors, 4);
+ reader.optimize_for_search();
+
const float query[] = {0.0f, 0.0f};
auto result = reader.search(query, 2, 2);
ASSERT_EQ(result.ids.size(), 2);
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 5e27ca8..5e58f92 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -526,6 +526,19 @@ pub unsafe extern "C" fn paimon_vindex_reader_metadata(
})
}
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_reader_optimize_for_search(
+ handle: *mut PaimonVindexReaderHandle,
+) -> c_int {
+ ffi_status(|| {
+ let handle = unsafe { reader_mut(handle) }?;
+ handle
+ .inner
+ .optimize_for_search()
+ .map_err(|e| format!("optimize_for_search: {}", e))
+ })
+}
+
#[no_mangle]
pub unsafe extern "C" fn paimon_vindex_reader_search(
handle: *mut PaimonVindexReaderHandle,
diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp
index ec91d32..b0ab679 100644
--- a/include/paimon_vindex.hpp
+++ b/include/paimon_vindex.hpp
@@ -250,6 +250,10 @@ public:
return result;
}
+ void optimize_for_search() {
+ check(paimon_vindex_reader_optimize_for_search(handle_));
+ }
+
SearchResult search(const float* query, size_t top_k, size_t nprobe,
size_t ef_search = 0) {
SearchResult result;
result.ids.resize(top_k);
diff --git
a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java
b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java
index 5350639..4adf3c9 100644
--- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java
+++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java
@@ -37,6 +37,8 @@ final class VectorIndexNative {
static native VectorIndexMetadata metadata(long ptr);
+ static native void optimizeForSearch(long ptr);
+
static native VectorSearchResult search(long ptr, float[] query, int k,
int nprobe, int efSearch);
static native VectorSearchResult searchWithRoaringFilter(
diff --git
a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java
b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java
index 34eefc7..2494cfc 100644
--- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java
+++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java
@@ -66,6 +66,17 @@ public final class VectorIndexReader implements
AutoCloseable {
return metadata().totalVectors();
}
+ public void optimizeForSearch() {
+ synchronized (nativeHandleLock) {
+ enterNativeHandle();
+ try {
+ VectorIndexNative.optimizeForSearch(requireOpen());
+ } finally {
+ exitNativeHandle();
+ }
+ }
+ }
+
public VectorSearchResult search(float[] query, int topK, int nprobe) {
return search(query, topK, nprobe, 0);
}
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 2cf1a05..6fd37a6 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
@@ -128,6 +128,12 @@ public class VectorIndexJavaApiTest {
reader.totalVectors();
}
});
+ assertThrows(IllegalStateException.class, new ThrowingRunnable() {
+ @Override
+ public void run() {
+ reader.optimizeForSearch();
+ }
+ });
assertThrows(IllegalStateException.class, new ThrowingRunnable() {
@Override
public void run() {
@@ -183,6 +189,7 @@ public class VectorIndexJavaApiTest {
reader.indexType();
reader.dimension();
reader.totalVectors();
+ reader.optimizeForSearch();
reader.search(new float[] {0.0f, 1.0f}, 10, 4);
reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32);
reader.search(new float[] {0.0f, 1.0f}, 10, 4, new byte[] {1, 2});
diff --git a/jni/src/lib.rs b/jni/src/lib.rs
index dd4aa8a..186fa84 100644
--- a/jni/src/lib.rs
+++ b/jni/src/lib.rs
@@ -494,6 +494,23 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_met
})
}
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_optimizeForSearch(
+ env: JNIEnv,
+ _class: JClass,
+ ptr: jlong,
+) {
+ jni_call_void(env, |env| {
+ let reader = match deref_reader(ptr) {
+ Some(reader) => reader,
+ None => return throw_and_return(env, "null native pointer (reader
already freed?)"),
+ };
+ if let Err(e) = reader.optimize_for_search() {
+ throw_and_return::<()>(env, &format!("optimize_for_search: {}",
e));
+ }
+ })
+}
+
#[no_mangle]
pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_search(
env: JNIEnv,
diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py
index 874567b..ddaedfa 100644
--- a/python/paimon_vindex/__init__.py
+++ b/python/paimon_vindex/__init__.py
@@ -298,6 +298,12 @@ class VectorIndexReader:
_check_error("metadata failed")
return _metadata_from_ffi(raw)
+ def optimize_for_search(self):
+ self._require_open()
+ rc = lib.paimon_vindex_reader_optimize_for_search(self._handle)
+ if rc != 0:
+ _check_error("optimize_for_search failed")
+
def _filter_args(self, filter_bytes):
if filter_bytes is None:
return None, 0, None
diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py
index 135a966..72cc25c 100644
--- a/python/paimon_vindex/_ffi.py
+++ b/python/paimon_vindex/_ffi.py
@@ -163,6 +163,9 @@ lib.paimon_vindex_reader_metadata.argtypes = [
]
lib.paimon_vindex_reader_metadata.restype = c_int
+lib.paimon_vindex_reader_optimize_for_search.argtypes = [c_void_p]
+lib.paimon_vindex_reader_optimize_for_search.restype = c_int
+
lib.paimon_vindex_reader_search.argtypes = [
c_void_p,
POINTER(c_float),
diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py
index c7e25c5..8380c84 100644
--- a/python/tests/test_vindex.py
+++ b/python/tests/test_vindex.py
@@ -108,9 +108,15 @@ def test_python_ffi_roundtrips_supported_indexes():
assert metadata.total_vectors == 512
ids, distances = reader.search(data[0], top_k=5, nprobe=4,
ef_search=32)
+ reader.optimize_for_search()
+ optimized_ids, optimized_distances = reader.search(
+ data[0], top_k=5, nprobe=4, ef_search=32
+ )
assert ids.shape == (5,)
assert distances.shape == (5,)
assert ids[0] == 0
+ np.testing.assert_array_equal(optimized_ids, ids)
+ np.testing.assert_allclose(optimized_distances, distances, rtol=0,
atol=1e-4)
def test_python_ffi_batch_search():
