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 47971a4 fix: align Python buffers before native calls (#113)
47971a4 is described below
commit 47971a458654366135c6252352f9c35b50a0ffa5
Author: Daoyuan Wang <[email protected]>
AuthorDate: Fri Oct 9 15:14:38 2026 +0800
fix: align Python buffers before native calls (#113)
---
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<usize, String> {
.ok_or_else(|| format!("{} length overflow", name))
}
+fn validate_slice_layout<T>(ptr: *const T, len: usize, name: &str) ->
Result<(), String> {
+ if len
+ .checked_mul(size_of::<T>())
+ .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::<u8>().add(1).cast::<f32>() };
+ 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::<f32>() + 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::<f32>(ptr::null(), 0, "query") }
+ .unwrap()
+ .is_empty());
+ assert!(unsafe { mut_slice::<f32>(ptr::null_mut(), 0, "output") }
+ .unwrap()
+ .is_empty());
+ }
+
struct BatchReadState {
data: Vec<u8>,
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()
