This is an automated email from the ASF dual-hosted git repository. jerry-024 pushed a commit to branch release-1-ali-4.1 in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
commit c8b64a96f568ceaea73c73382a7492853370fd30 Author: shyjsarah <[email protected]> AuthorDate: Fri Aug 7 10:32:41 2026 +0800 ffi: Add extensible C search parameters (#70) Co-authored-by: shaoyijie <[email protected]> (cherry picked from commit d5efa3f31116c347108e4c74dac804a75f169abb) --- c/test_vindex.c | 56 ++++++++ cpp/test_vindex.cpp | 15 ++ docs/api.html | 24 ++-- ffi/cbindgen.toml | 40 +++++- ffi/src/lib.rs | 359 ++++++++++++++++++++++++++++++++++++++++++++++ include/paimon_vindex.hpp | 33 +++-- 6 files changed, 507 insertions(+), 20 deletions(-) diff --git a/c/test_vindex.c b/c/test_vindex.c index ac8fd50..474aba7 100644 --- a/c/test_vindex.c +++ b/c/test_vindex.c @@ -57,6 +57,13 @@ enum { ROUNDTRIP_VECTOR_COUNT = ROUNDTRIP_NLIST * ROUNDTRIP_PER_LIST, }; +struct SearchParamsExPrefix { + uintptr_t struct_size; + uintptr_t top_k; + uint32_t search_width; + uintptr_t width; +}; + static void fail_ffi(const char *message) { const char *err = paimon_vindex_last_error(); fprintf(stderr, "%s: %s\n", message, err == NULL ? "(no error)" : err); @@ -351,6 +358,39 @@ static void run_roundtrip( } assert_id_in_cluster(batch_ids[0], 0); assert_id_in_cluster(batch_ids[1], 1); + struct PaimonVindexSearchParamsEx batch_params_ex = + paimon_vindex_search_params_ex_default(); + batch_params_ex.top_k = 1; + batch_params_ex.search_width = batch_params.search_width; + batch_params_ex.width = batch_params.width; + batch_params_ex.ivfpq_batch_table_reuse = + PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_OFF; + batch_params_ex.ivfpq_batch_table_reuse_max_bytes = 1; + if (paimon_vindex_reader_search_batch_ex( + reader, queries, 2, &batch_params_ex, batch_ids, batch_distances, 2) != 0) { + fail_ffi("reader search batch ex failed"); + } + assert_id_in_cluster(batch_ids[0], 0); + assert_id_in_cluster(batch_ids[1], 1); + struct SearchParamsExPrefix batch_params_prefix = { + .struct_size = + offsetof(struct SearchParamsExPrefix, width) + + sizeof(batch_params_prefix.width), + .top_k = 1, + .search_width = batch_params.search_width, + .width = batch_params.width}; + if (paimon_vindex_reader_search_batch_ex( + reader, + queries, + 2, + (const struct PaimonVindexSearchParamsEx *)&batch_params_prefix, + batch_ids, + batch_distances, + 2) != 0) { + fail_ffi("reader search batch ex prefix 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); @@ -496,7 +536,23 @@ static void test_supported_index_roundtrips(void) { 4); } +static void test_extensible_search_params_defaults(void) { + PaimonVindexSearchParamsEx params = paimon_vindex_search_params_ex_default(); + + ASSERT_TRUE( + params.struct_size == PAIMON_VINDEX_SEARCH_PARAMS_EX_V1_SIZE); + ASSERT_TRUE(params.search_width == PAIMON_VINDEX_SEARCH_WIDTH_AUTO); + ASSERT_TRUE( + params.ivfpq_batch_table_reuse == + PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO); + ASSERT_TRUE( + params.ivfpq_batch_table_reuse_max_bytes == + PAIMON_VINDEX_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES); + printf("PASS extensible_search_params_defaults\n"); +} + int main(void) { + test_extensible_search_params_defaults(); test_supported_index_roundtrips(); test_output_write_callback_error_propagates(); test_output_flush_callback_error_propagates(); diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp index 248c809..af7104a 100644 --- a/cpp/test_vindex.cpp +++ b/cpp/test_vindex.cpp @@ -325,8 +325,23 @@ static void test_worker_callback_reentry_is_rejected() { printf("PASS worker_callback_reentry_is_rejected\n"); } +static void test_extensible_search_params_forward_query_tuning() { + auto params = paimon::vindex::SearchParams::automatic(10); + params.max_initial_filter_expansion_factor = 4; + params.ivfpq_batch_table_reuse = PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_ON; + params.ivfpq_batch_table_reuse_max_bytes = 32 * 1024 * 1024; + + auto raw = params.to_ffi_ex(); + ASSERT_EQ(raw.struct_size, PAIMON_VINDEX_SEARCH_PARAMS_EX_V1_SIZE); + ASSERT_EQ(raw.max_initial_filter_expansion_factor, 4); + ASSERT_EQ(raw.ivfpq_batch_table_reuse, PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_ON); + ASSERT_EQ(raw.ivfpq_batch_table_reuse_max_bytes, 32 * 1024 * 1024); + printf("PASS extensible_search_params_forward_query_tuning\n"); +} + int main() { test_supported_index_roundtrips(); test_worker_callback_reentry_is_rejected(); + test_extensible_search_params_forward_query_tuning(); return 0; } diff --git a/docs/api.html b/docs/api.html index af7063d..080c0ab 100644 --- a/docs/api.html +++ b/docs/api.html @@ -56,9 +56,11 @@ <div class="table-wrap"><table><thead><tr><th>Parameter</th><th>Applies to</th><th>Description</th></tr></thead><tbody> <tr><td><code>top_k</code></td><td>All indexes</td><td>Number of nearest neighbors returned for each query.</td></tr> <tr><td><code>Auto</code> width</td><td>All indexes</td><td>IVF starts at <code>max(8, ceil(nlist/16))</code>, adds enough average-list capacity for at least <code>4 × top_k</code> candidates, scales for filter selectivity, and progressively doubles when a filtered result is short. DiskANN uses a calibrated width when present, otherwise <code>max(100, 2 × top_k)</code>.</td></tr> - <tr><td><code>max_initial_filter_expansion_factor</code></td><td>Automatic IVF search in Rust and Java</td><td>Optional cap on inverse-selectivity expansion of the initial <code>nprobe</code>. Unset preserves the current unlimited behavior. A value of 1 keeps the unfiltered automatic width. Lower factors reduce initial search work but may reduce recall compared with uncapped automatic search. Progressive expansion occurs only when fewer than <code>top_k</code> valid results a [...] + <tr><td><code>max_initial_filter_expansion_factor</code></td><td>Automatic IVF search in Rust, C, C++, and Java</td><td>Optional cap on inverse-selectivity expansion of the initial <code>nprobe</code>. Unset preserves the current unlimited behavior. A value of 1 keeps the unfiltered automatic width. Lower factors reduce initial search work but may reduce recall compared with uncapped automatic search. Progressive expansion occurs only when fewer than <code>top_k</code> valid [...] <tr><td><code>nprobe</code></td><td>IVF families</td><td>Explicit expert override for the number of lists probed. A tagged IVF width is rejected by DiskANN rather than silently ignored.</td></tr> <tr><td><code>l_search</code></td><td>DiskANN</td><td>Explicit expert override for graph search-list size; it is clamped to at least <code>top_k</code>. A tagged DiskANN width is rejected by IVF indexes.</td></tr> + <tr><td><code>ivfpq_batch_table_reuse</code></td><td>8-bit IVF-PQ batch search</td><td>Controls reuse of query distance tables across inverted lists: <code>off</code>, <code>on</code>, or <code>auto</code>. The default is <code>auto</code>.</td></tr> + <tr><td><code>ivfpq_batch_table_reuse_max_bytes</code></td><td>8-bit IVF-PQ batch search</td><td>Positive memory budget for reused query distance tables. The default is 512 MiB.</td></tr> </tbody></table></div> <p>Use the binding's <code>automatic</code> search-parameter constructor. For DiskANN, <code>calibrate_search_width</code> / <code>calibrateSearchWidth</code> evaluates widths 100, 200, and 400 on representative queries and remembers the smallest adjacent pair with at least 98% Top-K result overlap. This is a stability proxy, not a ground-truth recall guarantee; explicit widths always win.</p> <div class="callout"><strong>IVF-RQ width is a build option</strong><code>rq.bits</code> accepts 1–8 and defaults to 4. It is persisted in the file and reported as <code>rq_bits</code> / <code>rqBits</code> metadata; search has no separate bit-width switch.</div> @@ -167,15 +169,15 @@ paimon_vindex_reader_optimize_for_search(reader); int64_t ids[10]; float distances[10]; -PaimonVindexSearchParams params = { - .top_k = 10, - .search_width = PAIMON_VINDEX_SEARCH_WIDTH_AUTO, - .width = 0, -}; +PaimonVindexSearchParamsEx params = + paimon_vindex_search_params_ex_default(); +params.top_k = 10; +params.max_initial_filter_expansion_factor = 4; paimon_vindex_reader_warmup_queries(reader, representative_queries, 8, 0); -paimon_vindex_reader_search(reader, query, params, ids, distances, 10); +paimon_vindex_reader_search_ex(reader, query, &params, ids, distances, 10); paimon_vindex_reader_free(reader);</code></pre></div> - <p><code>PaimonVindexOutputFile</code> and <code>PaimonVindexInputFile</code> are callback structures. The input callback receives every positional range in one I/O batch, allowing object-store implementations to issue reads concurrently. DiskANN batch search may also invoke the callback concurrently from separate query workers, so the callback and its context must be thread-safe. The input structure also carries the three optional read-capability hints described above. C calle [...] + <p><code>PaimonVindexSearchParamsEx</code> is append-only and passed by pointer. Prefer <code>paimon_vindex_search_params_ex_default()</code>, which fills the semantic defaults and sets <code>struct_size</code> to <code>PAIMON_VINDEX_SEARCH_PARAMS_EX_V1_SIZE</code>. The size is the exact end of the last initialized field, not <code>sizeof(PaimonVindexSearchParamsEx)</code>, because structure sizes may include tail padding. Newer libraries default fields outside an older caller [...] + <p><code>PaimonVindexOutputFile</code> and <code>PaimonVindexInputFile</code> are callback structures. The input callback receives every positional range in one I/O batch, allowing object-store implementations to issue reads concurrently. DiskANN batch search may also invoke the callback concurrently from separate query workers, so the callback and its context must be thread-safe. The input structure also carries the three optional read-capability hints described above. C call [...] </section> <section class="article-section" id="cpp"> @@ -199,8 +201,10 @@ paimon::vindex::Reader reader(input_file); auto metadata = reader.metadata(); reader.optimize_for_search(); reader.warmup_queries(representative_queries.data(), 8); -auto result = reader.search( - query.data(), paimon::vindex::SearchParams::automatic(10));</code></pre></div> +auto params = paimon::vindex::SearchParams::automatic(10); +params.max_initial_filter_expansion_factor = 4; +params.ivfpq_batch_table_reuse = PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO; +auto result = reader.search(query.data(), params);</code></pre></div> </section> <section class="article-section" id="java"> diff --git a/ffi/cbindgen.toml b/ffi/cbindgen.toml index 0d64920..3a32e63 100644 --- a/ffi/cbindgen.toml +++ b/ffi/cbindgen.toml @@ -16,9 +16,45 @@ # under the License. language = "C" -header = "/* Generated by cbindgen - do not edit manually */" -include_guard = "PAIMON_VINDEX_H" +header = ''' +/* Generated by cbindgen - do not edit manually */ + +#ifndef PAIMON_VINDEX_H +#define PAIMON_VINDEX_H +''' autogen_warning = "/* Warning: this file is autogenerated by cbindgen. Don't modify this manually. */" +after_includes = ''' +#include <stddef.h> +''' +trailer = ''' +/** + * Exact initialized-prefix size for the first extended search parameter ABI. + * + * Do not replace this with sizeof(PaimonVindexSearchParamsEx): sizeof may + * include tail padding that a future ABI revision can reuse for another field. + */ +#define PAIMON_VINDEX_SEARCH_PARAMS_EX_V1_SIZE \ + (offsetof(PaimonVindexSearchParamsEx, ivfpq_batch_table_reuse_max_bytes) + \ + sizeof(((PaimonVindexSearchParamsEx *)0)->ivfpq_batch_table_reuse_max_bytes)) + +/** + * Returns V1 extended search parameters initialized with semantic defaults. + */ +static inline PaimonVindexSearchParamsEx paimon_vindex_search_params_ex_default(void) { + PaimonVindexSearchParamsEx params; + params.struct_size = PAIMON_VINDEX_SEARCH_PARAMS_EX_V1_SIZE; + params.top_k = 0; + params.search_width = PAIMON_VINDEX_SEARCH_WIDTH_AUTO; + params.width = 0; + params.max_initial_filter_expansion_factor = 0; + params.ivfpq_batch_table_reuse = PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO; + params.ivfpq_batch_table_reuse_max_bytes = + PAIMON_VINDEX_DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES; + return params; +} + +#endif /* PAIMON_VINDEX_H */ +''' tab_width = 4 style = "both" diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 31bcc36..90b9916 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -29,6 +29,7 @@ use std::cell::RefCell; use std::collections::HashMap; use std::ffi::{CStr, CString}; use std::io; +use std::mem::{offset_of, size_of}; use std::os::raw::{c_char, c_int, c_void}; use std::panic::{self, AssertUnwindSafe}; use std::{ptr, slice}; @@ -282,6 +283,25 @@ pub struct PaimonVindexSearchParamsV2 { pub ivfpq_batch_table_reuse_max_bytes: usize, } +/// Extensible search parameters passed by pointer. +/// +/// Callers must set `struct_size` to the exact end of the last initialized +/// field, not `size_of` the allocation. Future versions may append fields; +/// readers use `struct_size` to default fields that are not present and ignore +/// unknown trailing fields. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct PaimonVindexSearchParamsEx { + pub struct_size: usize, + pub top_k: usize, + pub search_width: u32, + pub width: usize, + /// Zero leaves the automatic IVF filter expansion uncapped. + pub max_initial_filter_expansion_factor: usize, + pub ivfpq_batch_table_reuse: u32, + pub ivfpq_batch_table_reuse_max_bytes: usize, +} + #[repr(C)] #[derive(Clone, Copy)] pub struct PaimonVindexReaderOptions { @@ -568,6 +588,127 @@ fn search_params_v2_from_ffi( Ok(result) } +/// Reads an optional field from an append-only FFI structure. +/// +/// # Safety +/// +/// `params` must point to at least `struct_size` readable bytes. `offset` must +/// identify a field with C layout and type `T` in that allocation. +unsafe fn search_params_ex_field<T: Copy>( + params: *const PaimonVindexSearchParamsEx, + struct_size: usize, + offset: usize, + default: T, +) -> T { + let Some(field_end) = offset.checked_add(size_of::<T>()) else { + return default; + }; + if field_end > struct_size { + return default; + } + // SAFETY: The caller guarantees that `params` covers `struct_size` + // readable bytes, and the bounds check above proves that this field lies + // within that region. Unaligned reads support C layouts with padding. + unsafe { ptr::read_unaligned(params.cast::<u8>().add(offset).cast::<T>()) } +} + +/// Converts an append-only C search-parameter structure. +/// +/// # Safety +/// +/// `params` must be null or point to a readable allocation whose first +/// `usize` contains its actual byte size. +unsafe fn search_params_ex_from_ffi( + params: *const PaimonVindexSearchParamsEx, +) -> Result<VectorSearchParams, String> { + if params.is_null() { + return Err("search params pointer is null".to_string()); + } + // SAFETY: The function contract requires the first `usize` to be readable. + let struct_size = unsafe { ptr::read_unaligned(params.cast::<usize>()) }; + let required_size = offset_of!(PaimonVindexSearchParamsEx, width) + size_of::<usize>(); + if struct_size < required_size { + return Err(format!( + "search params struct size {} is smaller than required {}", + struct_size, required_size + )); + } + + let top_k = unsafe { + search_params_ex_field( + params, + struct_size, + offset_of!(PaimonVindexSearchParamsEx, top_k), + 0, + ) + }; + let search_width = unsafe { + search_params_ex_field( + params, + struct_size, + offset_of!(PaimonVindexSearchParamsEx, search_width), + PAIMON_VINDEX_SEARCH_WIDTH_AUTO, + ) + }; + let width = unsafe { + search_params_ex_field( + params, + struct_size, + offset_of!(PaimonVindexSearchParamsEx, width), + 0, + ) + }; + let mut result = search_params_from_ffi(PaimonVindexSearchParams { + top_k, + search_width, + width, + })?; + + let max_initial_filter_expansion_factor = unsafe { + search_params_ex_field( + params, + struct_size, + offset_of!( + PaimonVindexSearchParamsEx, + max_initial_filter_expansion_factor + ), + 0, + ) + }; + result.max_initial_filter_expansion_factor = + (max_initial_filter_expansion_factor != 0).then_some(max_initial_filter_expansion_factor); + + let reuse_mode = unsafe { + search_params_ex_field( + params, + struct_size, + offset_of!(PaimonVindexSearchParamsEx, ivfpq_batch_table_reuse), + PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO, + ) + }; + result.ivfpq_batch_table_reuse = match reuse_mode { + 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}")), + }; + result.ivfpq_batch_table_reuse_max_bytes = unsafe { + search_params_ex_field( + params, + struct_size, + offset_of!( + PaimonVindexSearchParamsEx, + ivfpq_batch_table_reuse_max_bytes + ), + DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES, + ) + }; + if result.ivfpq_batch_table_reuse_max_bytes == 0 { + return Err("IVF-PQ batch table reuse max bytes must be positive".to_string()); + } + Ok(result) +} + // ======================== Trainer / Writer ======================== #[no_mangle] @@ -917,6 +1058,34 @@ pub unsafe extern "C" fn paimon_vindex_reader_search( }) } +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_ex( + handle: *mut PaimonVindexReaderHandle, + query: *const f32, + params: *const PaimonVindexSearchParamsEx, + out_ids: *mut i64, + out_distances: *mut f32, + result_len: usize, +) -> c_int { + ffi_status(|| { + let handle = unsafe { reader_mut(handle) }?; + let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; + let params = unsafe { search_params_ex_from_ffi(params) }?; + let (ids, distances) = handle + .inner + .search(query, params) + .map_err(|e| format!("search: {}", e))?; + copy_search_result( + &ids, + &distances, + out_ids, + out_distances, + result_len, + params.top_k, + ) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter( handle: *mut PaimonVindexReaderHandle, @@ -948,6 +1117,37 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter( }) } +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter_ex( + handle: *mut PaimonVindexReaderHandle, + query: *const f32, + params: *const PaimonVindexSearchParamsEx, + 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 = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; + let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; + let params = unsafe { search_params_ex_from_ffi(params) }?; + let (ids, distances) = handle + .inner + .search_with_roaring_filter(query, params, filter) + .map_err(|e| format!("search_with_roaring_filter: {}", e))?; + copy_search_result( + &ids, + &distances, + out_ids, + out_distances, + result_len, + params.top_k, + ) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_reader_search_batch( handle: *mut PaimonVindexReaderHandle, @@ -979,6 +1179,37 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch( }) } +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_batch_ex( + handle: *mut PaimonVindexReaderHandle, + queries: *const f32, + query_count: usize, + params: *const PaimonVindexSearchParamsEx, + 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 = unsafe { search_params_ex_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_v2( handle: *mut PaimonVindexReaderHandle, @@ -1044,6 +1275,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_ex( + handle: *mut PaimonVindexReaderHandle, + queries: *const f32, + query_count: usize, + params: *const PaimonVindexSearchParamsEx, + 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 = unsafe { search_params_ex_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, + ) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter_v2( handle: *mut PaimonVindexReaderHandle, @@ -1237,4 +1502,98 @@ mod tests { .is_err()); } } + + #[test] + fn ffi_extended_search_parameters_preserve_all_query_tuning_options() { + let raw = PaimonVindexSearchParamsEx { + struct_size: size_of::<PaimonVindexSearchParamsEx>(), + top_k: 10, + search_width: PAIMON_VINDEX_SEARCH_WIDTH_AUTO, + width: 0, + max_initial_filter_expansion_factor: 4, + ivfpq_batch_table_reuse: PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_ON, + ivfpq_batch_table_reuse_max_bytes: 32 * 1024 * 1024, + }; + + let params = unsafe { search_params_ex_from_ffi(&raw) }.unwrap(); + + assert_eq!(params.max_initial_filter_expansion_factor, Some(4)); + assert_eq!(params.ivfpq_batch_table_reuse, IvfPqBatchTableReuseMode::On); + assert_eq!(params.ivfpq_batch_table_reuse_max_bytes, 32 * 1024 * 1024); + } + + #[repr(C)] + struct SearchParamsExPrefix { + struct_size: usize, + top_k: usize, + search_width: u32, + width: usize, + } + + #[test] + fn ffi_extended_search_parameters_default_fields_missing_from_shorter_structs() { + let raw = SearchParamsExPrefix { + struct_size: size_of::<SearchParamsExPrefix>(), + top_k: 10, + search_width: PAIMON_VINDEX_SEARCH_WIDTH_IVF_NPROBE, + width: 16, + }; + + let params = unsafe { + search_params_ex_from_ffi( + (&raw as *const SearchParamsExPrefix).cast::<PaimonVindexSearchParamsEx>(), + ) + } + .unwrap(); + + assert_eq!(params.max_initial_filter_expansion_factor, None); + assert_eq!( + 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_extended_search_parameters_reject_undersized_structs() { + let raw_size = size_of::<usize>(); + let error = unsafe { + search_params_ex_from_ffi( + (&raw_size as *const usize).cast::<PaimonVindexSearchParamsEx>(), + ) + } + .unwrap_err(); + + assert!(error.contains("smaller than required")); + } + + #[repr(C)] + struct FutureSearchParamsEx { + current: PaimonVindexSearchParamsEx, + future_field: u64, + } + + #[test] + fn ffi_extended_search_parameters_ignore_unknown_trailing_fields() { + let raw = FutureSearchParamsEx { + current: PaimonVindexSearchParamsEx { + struct_size: size_of::<FutureSearchParamsEx>(), + top_k: 10, + search_width: PAIMON_VINDEX_SEARCH_WIDTH_AUTO, + width: 0, + max_initial_filter_expansion_factor: 2, + ivfpq_batch_table_reuse: PAIMON_VINDEX_IVFPQ_BATCH_TABLE_REUSE_AUTO, + ivfpq_batch_table_reuse_max_bytes: DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES, + }, + future_field: u64::MAX, + }; + + let params = unsafe { search_params_ex_from_ffi(&raw.current) }.unwrap(); + + assert_eq!(params.top_k, 10); + assert_eq!(params.max_initial_filter_expansion_factor, Some(2)); + } } diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp index 0b3e8e6..f3638bb 100644 --- a/include/paimon_vindex.hpp +++ b/include/paimon_vindex.hpp @@ -228,6 +228,7 @@ struct SearchParams { size_t top_k = 0; uint32_t search_width = PAIMON_VINDEX_SEARCH_WIDTH_AUTO; size_t width = 0; + size_t max_initial_filter_expansion_factor = 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; @@ -269,6 +270,18 @@ struct SearchParams { return params; } + PaimonVindexSearchParamsEx to_ffi_ex() const { + auto params = paimon_vindex_search_params_ex_default(); + params.top_k = top_k; + params.search_width = search_width; + params.width = width; + params.max_initial_filter_expansion_factor = + max_initial_filter_expansion_factor; + 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; }; @@ -576,10 +589,11 @@ public: SearchResult result; result.ids.resize(params.top_k); result.distances.resize(params.top_k); - check(paimon_vindex_reader_search( + auto raw_params = params.to_ffi_ex(); + check(paimon_vindex_reader_search_ex( require_open(), query, - params.to_ffi(), + &raw_params, result.ids.data(), result.distances.data(), params.top_k)); @@ -595,10 +609,11 @@ public: SearchResult result; result.ids.resize(params.top_k); result.distances.resize(params.top_k); - check(paimon_vindex_reader_search_with_roaring_filter( + auto raw_params = params.to_ffi_ex(); + check(paimon_vindex_reader_search_with_roaring_filter_ex( require_open(), query, - params.to_ffi(), + &raw_params, filter, filter_len, result.ids.data(), @@ -616,11 +631,12 @@ public: SearchResult result; result.ids.resize(result_len); result.distances.resize(result_len); - check(paimon_vindex_reader_search_batch_v2( + auto raw_params = params.to_ffi_ex(); + check(paimon_vindex_reader_search_batch_ex( require_open(), queries, query_count, - params.to_ffi_v2(), + &raw_params, result.ids.data(), result.distances.data(), result_len)); @@ -638,11 +654,12 @@ public: SearchResult result; result.ids.resize(result_len); result.distances.resize(result_len); - check(paimon_vindex_reader_search_batch_with_roaring_filter_v2( + auto raw_params = params.to_ffi_ex(); + check(paimon_vindex_reader_search_batch_with_roaring_filter_ex( require_open(), queries, query_count, - params.to_ffi_v2(), + &raw_params, filter, filter_len, result.ids.data(),
