From a7d4c6a630f60876f297b0563234f051a810410f Mon Sep 17 00:00:00 2001 From: Daoyuan Wang Date: Fri, 9 Oct 2026 11:13:37 +0800 Subject: [PATCH] fix: align Python buffers before native calls Require aligned NumPy buffers for training, warm-up, and Top-K operations. Reject misaligned and oversized FFI slices before constructing Rust references. Validation: 18 FFI tests; 405 Python tests passed, with one optional range-fixture test skipped; FFI Clippy passed. --- ffi/src/lib.rs | 42 ++++++++++++++++++++++++++++++++ python/paimon_vindex/__init__.py | 6 ++--- python/tests/test_vindex.py | 32 ++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 3 deletions(-) diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index b6dc138..ba831fb 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -504,12 +504,26 @@ fn checked_len(a: usize, b: usize, name: &str) -> Result { .ok_or_else(|| format!("{} length overflow", name)) } +fn validate_slice_layout(ptr: *const T, len: usize, name: &str) -> Result<(), String> { + if len + .checked_mul(size_of::()) + .is_none_or(|bytes| bytes > isize::MAX as usize) + { + return Err(format!("{name} byte length overflow")); + } + if !ptr.is_aligned() { + return Err(format!("{name} pointer is not aligned")); + } + Ok(()) +} + unsafe fn const_slice<'a, T>(ptr: *const T, len: usize, name: &str) -> Result<&'a [T], String> { if len == 0 { Ok(&[]) } else if ptr.is_null() { Err(format!("{} pointer is null", name)) } else { + validate_slice_layout(ptr, len, name)?; Ok(unsafe { slice::from_raw_parts(ptr, len) }) } } @@ -520,6 +534,7 @@ unsafe fn mut_slice<'a, T>(ptr: *mut T, len: usize, name: &str) -> Result<&'a mu } else if ptr.is_null() { Err(format!("{} pointer is null", name)) } else { + validate_slice_layout(ptr, len, name)?; Ok(unsafe { slice::from_raw_parts_mut(ptr, len) }) } } @@ -1353,6 +1368,33 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter_v mod tests { use super::*; + #[test] + fn ffi_slices_reject_unaligned_and_oversized_buffers() { + let mut storage = [0u32; 2]; + let unaligned = unsafe { storage.as_mut_ptr().cast::().add(1).cast::() }; + assert!(unsafe { const_slice(unaligned, 1, "query") } + .unwrap_err() + .contains("not aligned")); + assert!(unsafe { mut_slice(unaligned, 1, "output") } + .unwrap_err() + .contains("not aligned")); + let oversized = isize::MAX as usize / size_of::() + 1; + assert!(unsafe { const_slice(storage.as_ptr(), oversized, "query") } + .unwrap_err() + .contains("overflow")); + assert!( + unsafe { mut_slice(storage.as_mut_ptr(), oversized, "output") } + .unwrap_err() + .contains("overflow") + ); + assert!(unsafe { const_slice::(ptr::null(), 0, "query") } + .unwrap() + .is_empty()); + assert!(unsafe { mut_slice::(ptr::null_mut(), 0, "output") } + .unwrap() + .is_empty()); + } + struct BatchReadState { data: Vec, calls: usize, diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py index 0fda4b1..e817349 100644 --- a/python/paimon_vindex/__init__.py +++ b/python/paimon_vindex/__init__.py @@ -464,14 +464,14 @@ def _float32_matrix(value, name): array = np.asarray(value, dtype=np.float32) if array.ndim != 2: raise ValueError(f"{name} must be a two-dimensional float32 array") - return np.ascontiguousarray(array) + return np.require(array, dtype=np.float32, requirements=["C", "A"]) def _float32_vector(value, name): array = np.asarray(value, dtype=np.float32) if array.ndim != 1: raise ValueError(f"{name} must be a one-dimensional float32 array") - return np.ascontiguousarray(array) + return np.require(array, dtype=np.float32, requirements=["C", "A"]) def _require_range_search(): @@ -551,7 +551,7 @@ def _int64_vector(value, name): array = np.asarray(value, dtype=np.int64) if array.ndim != 1: raise ValueError(f"{name} must be a one-dimensional int64 array") - return np.ascontiguousarray(array) + return np.require(array, dtype=np.int64, requirements=["C", "A"]) def _bytes_buffer(value, name): diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py index d82490b..19ffdd0 100644 --- a/python/tests/test_vindex.py +++ b/python/tests/test_vindex.py @@ -63,6 +63,38 @@ def reader_from_bytes(data): return VectorIndexReader(VectorIndexInput(data)) +def test_python_native_operations_accept_unaligned_numpy_buffers(): + def unaligned_copy(source): + result = np.ndarray( + source.shape, dtype=source.dtype, + buffer=bytearray(source.nbytes + 1), offset=1, + ) + result[:] = source + assert result.flags.c_contiguous and not result.flags.aligned + return result + + data = clustered_data(512, 8, 4) + ids = np.arange(len(data), dtype=np.int64) + options = {"index.type": "ivf_flat", "dimension": "8", "nlist": "4", "metric": "l2"} + output = io.BytesIO() + with VectorIndexTrainer.train(options, unaligned_copy(data)) as training: + with VectorIndexWriter(training) as writer: + writer.add_vectors(unaligned_copy(ids), unaligned_copy(data)) + writer.write(output) + + with reader_from_bytes(output.getvalue()) as reader: + reader.warmup_queries(unaligned_copy(data[:2])) + params = SearchParams.ivf(5, 4) + for actual, expected in zip( + reader.search(unaligned_copy(data[0]), params), reader.search(data[0], params) + ): + np.testing.assert_array_equal(actual, expected) + for actual, expected in zip( + reader.search_batch(unaligned_copy(data[:2]), params), reader.search_batch(data[:2], params) + ): + np.testing.assert_array_equal(actual, expected) + + def test_python_search_parameters_remain_algorithm_specific(): params = SearchParams.diskann(top_k=10, l_search=200).to_ffi()