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 1499263 Batch positional reads across C and Python (#60)
1499263 is described below
commit 1499263aeadfb3f3d94111c2718e8a22bc7dc443
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Jul 22 15:51:27 2026 +0800
Batch positional reads across C and Python (#60)
---
README.md | 6 +-
c/test_vindex.c | 52 ++++++++++------
cpp/test_vindex.cpp | 16 ++++-
ffi/src/lib.rs | 127 +++++++++++++++++++++++++++++++++------
include/paimon_vindex.hpp | 13 ++--
python/paimon_vindex/__init__.py | 41 ++++++++-----
python/paimon_vindex/_ffi.py | 19 +++++-
python/tests/test_vindex.py | 29 +++++++++
tools/README.md | 19 ++++++
tools/check_license_headers.py | 1 +
10 files changed, 259 insertions(+), 64 deletions(-)
diff --git a/README.md b/README.md
index a949786..4129065 100644
--- a/README.md
+++ b/README.md
@@ -197,8 +197,10 @@ paimon_vindex_reader_search(
paimon_vindex_reader_free(reader);
```
-`PaimonVindexOutputFile` and `PaimonVindexInputFile` are callback structs. C
-callers own result buffers and pass their capacity to search calls. Functions
+`PaimonVindexOutputFile` and `PaimonVindexInputFile` are callback structs. The
+input callback receives all positional read ranges for one I/O batch in a
+single call, so object-store implementations can issue the reads concurrently.
+C callers own result buffers and pass their capacity to search calls. Functions
return `0` on success and `-1` on error; `paimon_vindex_last_error()` returns a
thread-local error string.
diff --git a/c/test_vindex.c b/c/test_vindex.c
index 7fa3f90..1cb05cb 100644
--- a/c/test_vindex.c
+++ b/c/test_vindex.c
@@ -47,6 +47,7 @@ struct MemBuffer {
size_t len;
size_t cap;
size_t pos;
+ size_t max_read_request_count;
};
enum {
@@ -113,17 +114,27 @@ static int64_t mem_pos(void *ctx) {
return (int64_t)buf->pos;
}
-static int mem_read_at(void *ctx, uint64_t offset, uint8_t *dst, uintptr_t
len) {
+static int mem_read_ranges(
+ void *ctx,
+ struct PaimonVindexReadRequest *requests,
+ uintptr_t request_count) {
struct MemBuffer *buf = (struct MemBuffer *)ctx;
- if (offset > SIZE_MAX || len > SIZE_MAX) {
- return -1;
+ if (request_count > buf->max_read_request_count) {
+ buf->max_read_request_count = (size_t)request_count;
}
- size_t off = (size_t)offset;
- size_t n = (size_t)len;
- if (off > buf->len || n > buf->len - off) {
- return -1;
+ for (uintptr_t i = 0; i < request_count; i++) {
+ uint64_t offset = requests[i].offset;
+ uintptr_t len = requests[i].len;
+ if (offset > SIZE_MAX || len > SIZE_MAX) {
+ return -1;
+ }
+ size_t off = (size_t)offset;
+ size_t n = (size_t)len;
+ if (off > buf->len || n > buf->len - off) {
+ return -1;
+ }
+ memcpy(requests[i].buf, buf->data + off, n);
}
- memcpy(dst, buf->data + off, n);
return 0;
}
@@ -139,11 +150,13 @@ static int failing_flush(void *ctx) {
return -1;
}
-static int failing_read_at(void *ctx, uint64_t offset, uint8_t *dst, uintptr_t
len) {
+static int failing_read_ranges(
+ void *ctx,
+ struct PaimonVindexReadRequest *requests,
+ uintptr_t request_count) {
(void)ctx;
- (void)offset;
- (void)dst;
- (void)len;
+ (void)requests;
+ (void)request_count;
return -1;
}
@@ -235,7 +248,7 @@ static void run_roundtrip(
struct PaimonVindexInputFile input = {
.ctx = &buf,
- .read_at_fn = mem_read_at,
+ .read_ranges_fn = mem_read_ranges,
};
PaimonVindexReaderHandle *reader = paimon_vindex_reader_open(input);
if (reader == NULL) {
@@ -269,6 +282,9 @@ static void run_roundtrip(
}
assert_id_in_cluster(result_ids[0], 0);
ASSERT_TRUE(isfinite(result_distances[0]));
+ if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_PQ) {
+ ASSERT_TRUE(buf.max_read_request_count > 1);
+ }
if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) {
search_params.query_bits = 4;
if (paimon_vindex_reader_search(
@@ -370,16 +386,16 @@ static void
test_output_flush_callback_error_propagates(void) {
printf("PASS output_flush_callback_error_propagates\n");
}
-static void test_input_read_callback_error_propagates(void) {
+static void test_input_read_ranges_callback_error_propagates(void) {
struct PaimonVindexInputFile input = {
.ctx = NULL,
- .read_at_fn = failing_read_at,
+ .read_ranges_fn = failing_read_ranges,
};
PaimonVindexReaderHandle *reader = paimon_vindex_reader_open(input);
ASSERT_TRUE(reader == NULL);
- assert_last_error_contains("read_at callback failed");
- printf("PASS input_read_callback_error_propagates\n");
+ assert_last_error_contains("read_ranges callback failed");
+ printf("PASS input_read_ranges_callback_error_propagates\n");
}
static void test_supported_index_roundtrips(void) {
@@ -443,6 +459,6 @@ int main(void) {
test_supported_index_roundtrips();
test_output_write_callback_error_propagates();
test_output_flush_callback_error_propagates();
- test_input_read_callback_error_propagates();
+ test_input_read_ranges_callback_error_propagates();
return 0;
}
diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp
index be842ae..9cc31d6 100644
--- a/cpp/test_vindex.cpp
+++ b/cpp/test_vindex.cpp
@@ -43,6 +43,7 @@
struct MemBuffer {
std::vector<uint8_t> data;
size_t pos = 0;
+ mutable size_t max_read_request_count = 0;
};
constexpr size_t kRoundtripDimension = 8;
@@ -64,9 +65,15 @@ static paimon::vindex::OutputFile make_output(MemBuffer&
buf) {
static paimon::vindex::InputFile make_input(const MemBuffer& buf) {
paimon::vindex::InputFile in;
- in.read_at_fn = [&buf](uint64_t offset, uint8_t* dst, size_t len) -> int {
- if (offset + len > buf.data.size()) return -1;
- memcpy(dst, buf.data.data() + offset, len);
+ in.read_ranges_fn = [&buf](
+ paimon::vindex::ReadRequest* requests,
+ size_t request_count) -> int {
+ buf.max_read_request_count = std::max(buf.max_read_request_count,
request_count);
+ for (size_t i = 0; i < request_count; i++) {
+ const auto& request = requests[i];
+ if (request.offset + request.len > buf.data.size()) return -1;
+ memcpy(request.buf, buf.data.data() + request.offset, request.len);
+ }
return 0;
};
return in;
@@ -153,6 +160,9 @@ static void run_roundtrip(
ASSERT_EQ(result.ids.size(), 2);
assert_id_in_cluster(result.ids[0], 0);
ASSERT_TRUE(std::isfinite(result.distances[0]));
+ if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_PQ) {
+ ASSERT_TRUE(buf.max_read_request_count > 1);
+ }
if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) {
auto query_bits_result =
reader.search(query.data(), paimon::vindex::SearchParams{2, 4, 16,
4});
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 46ae9c8..6e9eb94 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -161,10 +161,25 @@ impl SeekWrite for FfiOutputFile {
}
}
+#[repr(C)]
+pub struct PaimonVindexReadRequest {
+ /// Absolute byte offset in the input file.
+ pub offset: u64,
+ /// Destination buffer that the callback must fill before returning.
+ pub buf: *mut u8,
+ /// Destination buffer length in bytes.
+ pub len: usize,
+}
+
#[repr(C)]
pub struct PaimonVindexInputFile {
pub ctx: *mut c_void,
- pub read_at_fn: Option<unsafe extern "C" fn(*mut c_void, u64, *mut u8,
usize) -> c_int>,
+ /// Reads every request in the batch, preferably concurrently.
+ ///
+ /// Request descriptors and their buffers are valid only for the duration
of
+ /// the callback and must not be retained by the implementation.
+ pub read_ranges_fn:
+ Option<unsafe extern "C" fn(*mut c_void, *mut PaimonVindexReadRequest,
usize) -> c_int>,
}
struct FfiInputFile {
@@ -175,27 +190,26 @@ unsafe impl Send for FfiInputFile {}
impl SeekRead for FfiInputFile {
fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> {
- if let Some(read_at_fn) = self.raw.read_at_fn {
- for range in ranges {
- let result = unsafe {
- read_at_fn(
- self.raw.ctx,
- range.pos,
- range.buf.as_mut_ptr(),
- range.buf.len(),
- )
- };
- if result != 0 {
- return Err(io::Error::other(format!(
- "read_at callback failed at offset {} length {}",
- range.pos,
- range.buf.len()
- )));
- }
+ if ranges.is_empty() {
+ return Ok(());
+ }
+ if let Some(read_ranges_fn) = self.raw.read_ranges_fn {
+ let mut requests = ranges
+ .iter_mut()
+ .map(|range| PaimonVindexReadRequest {
+ offset: range.pos,
+ buf: range.buf.as_mut_ptr(),
+ len: range.buf.len(),
+ })
+ .collect::<Vec<_>>();
+ let result =
+ unsafe { read_ranges_fn(self.raw.ctx, requests.as_mut_ptr(),
requests.len()) };
+ if result != 0 {
+ return Err(io::Error::other("read_ranges callback failed"));
}
Ok(())
} else {
- Err(io::Error::other("read_at_fn is null"))
+ Err(io::Error::other("read_ranges_fn is null"))
}
}
}
@@ -802,3 +816,78 @@ pub unsafe extern "C" fn
paimon_vindex_reader_search_batch_with_roaring_filter(
)
})
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ struct BatchReadState {
+ data: Vec<u8>,
+ calls: usize,
+ range_count: usize,
+ }
+
+ unsafe extern "C" fn read_ranges(
+ ctx: *mut c_void,
+ requests: *mut PaimonVindexReadRequest,
+ request_count: usize,
+ ) -> c_int {
+ let state = unsafe { &mut *(ctx as *mut BatchReadState) };
+ state.calls += 1;
+ state.range_count = request_count;
+ let requests = unsafe { std::slice::from_raw_parts_mut(requests,
request_count) };
+ for request in requests {
+ let start = request.offset as usize;
+ let end = start + request.len;
+ let destination = unsafe {
std::slice::from_raw_parts_mut(request.buf, request.len) };
+ destination.copy_from_slice(&state.data[start..end]);
+ }
+ 0
+ }
+
+ #[test]
+ fn ffi_pread_forwards_all_ranges_in_one_callback() {
+ let mut state = BatchReadState {
+ data: (0u8..32).collect(),
+ calls: 0,
+ range_count: 0,
+ };
+ let raw = PaimonVindexInputFile {
+ ctx: (&mut state as *mut BatchReadState).cast(),
+ read_ranges_fn: Some(read_ranges),
+ };
+ let mut input = FfiInputFile { raw };
+ let mut first = [0u8; 3];
+ let mut second = [0u8; 4];
+
+ input
+ .pread(&mut [
+ ReadRequest::new(2, &mut first),
+ ReadRequest::new(11, &mut second),
+ ])
+ .unwrap();
+
+ assert_eq!(state.calls, 1);
+ assert_eq!(state.range_count, 2);
+ assert_eq!(first, [2, 3, 4]);
+ assert_eq!(second, [11, 12, 13, 14]);
+ }
+
+ #[test]
+ fn ffi_pread_skips_callback_for_empty_ranges() {
+ let mut state = BatchReadState {
+ data: Vec::new(),
+ calls: 0,
+ range_count: 0,
+ };
+ let raw = PaimonVindexInputFile {
+ ctx: (&mut state as *mut BatchReadState).cast(),
+ read_ranges_fn: Some(read_ranges),
+ };
+ let mut input = FfiInputFile { raw };
+
+ input.pread(&mut []).unwrap();
+
+ assert_eq!(state.calls, 0);
+ }
+}
diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp
index 8584f09..294e482 100644
--- a/include/paimon_vindex.hpp
+++ b/include/paimon_vindex.hpp
@@ -51,8 +51,10 @@ struct OutputFile {
std::function<int64_t()> get_pos_fn;
};
+using ReadRequest = PaimonVindexReadRequest;
+
struct InputFile {
- std::function<int(uint64_t offset, uint8_t* buf, size_t len)> read_at_fn;
+ std::function<int(ReadRequest* requests, size_t request_count)>
read_ranges_fn;
};
namespace detail {
@@ -86,10 +88,13 @@ inline int64_t stream_get_pos(void* ctx) noexcept {
}
}
-inline int input_read_at(void* ctx, uint64_t offset, uint8_t* buf, size_t len)
noexcept {
+inline int input_read_ranges(
+ void* ctx,
+ PaimonVindexReadRequest* raw_requests,
+ size_t request_count) noexcept {
try {
auto* cbs = static_cast<InputFile*>(ctx);
- return cbs->read_at_fn(offset, buf, len);
+ return cbs->read_ranges_fn(raw_requests, request_count);
} catch (...) {
return -1;
}
@@ -326,7 +331,7 @@ public:
explicit Reader(InputFile input) :
input_(std::make_shared<InputFile>(std::move(input))) {
PaimonVindexInputFile raw;
raw.ctx = input_.get();
- raw.read_at_fn = detail::input_read_at;
+ raw.read_ranges_fn = detail::input_read_ranges;
handle_ = paimon_vindex_reader_open(raw);
if (!handle_) throw Error("failed to open vector index reader");
}
diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py
index 4bdd4a3..2113220 100644
--- a/python/paimon_vindex/__init__.py
+++ b/python/paimon_vindex/__init__.py
@@ -136,6 +136,29 @@ def _option_arrays(options: Mapping[str, str]):
return option_items, key_bytes, value_bytes, keys, values
+def _make_read_ranges_callback(input):
+ @_ffi.READ_RANGES_FN
+ def read_ranges_callback(ctx, requests, request_count):
+ try:
+ ranges = [
+ (requests[i].offset, requests[i].len)
+ for i in range(request_count)
+ ]
+ chunks = input.pread_many(ranges)
+ if len(chunks) != request_count:
+ return -1
+ for i, chunk in enumerate(chunks):
+ data = bytes(chunk)
+ if len(data) != requests[i].len:
+ return -1
+ ctypes.memmove(requests[i].buf, data, len(data))
+ return 0
+ except Exception:
+ return -1
+
+ return read_ranges_callback
+
+
class VectorIndexTraining:
def __init__(self, handle):
self._closed = False
@@ -382,24 +405,10 @@ class VectorIndexReader:
self._input = input
self._closed = False
- @_ffi.READ_AT_FN
- def read_at_callback(ctx, offset, buf, length):
- try:
- chunks = self._input.pread_many([(offset, length)])
- if len(chunks) != 1:
- return -1
- data = bytes(chunks[0])
- if len(data) != length:
- return -1
- ctypes.memmove(buf, data, length)
- return 0
- except Exception:
- return -1
-
- self._read_at_callback = read_at_callback
+ self._read_ranges_callback = _make_read_ranges_callback(self._input)
input_file = _ffi.PaimonVindexInputFile()
input_file.ctx = None
- input_file.read_at_fn = self._read_at_callback
+ input_file.read_ranges_fn = self._read_ranges_callback
self._handle = lib.paimon_vindex_reader_open(input_file)
if not self._handle:
_check_error("failed to open reader")
diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py
index c084a2e..fbd87fc 100644
--- a/python/paimon_vindex/_ffi.py
+++ b/python/paimon_vindex/_ffi.py
@@ -85,7 +85,22 @@ lib = _load_library()
WRITE_FN = CFUNCTYPE(c_int, c_void_p, POINTER(c_uint8), c_size_t)
FLUSH_FN = CFUNCTYPE(c_int, c_void_p)
GET_POS_FN = CFUNCTYPE(c_int64, c_void_p)
-READ_AT_FN = CFUNCTYPE(c_int, c_void_p, c_uint64, POINTER(c_uint8), c_size_t)
+
+
+class PaimonVindexReadRequest(Structure):
+ _fields_ = [
+ ("offset", c_uint64),
+ ("buf", POINTER(c_uint8)),
+ ("len", c_size_t),
+ ]
+
+
+READ_RANGES_FN = CFUNCTYPE(
+ c_int,
+ c_void_p,
+ POINTER(PaimonVindexReadRequest),
+ c_size_t,
+)
class PaimonVindexOutputFile(Structure):
@@ -100,7 +115,7 @@ class PaimonVindexOutputFile(Structure):
class PaimonVindexInputFile(Structure):
_fields_ = [
("ctx", c_void_p),
- ("read_at_fn", READ_AT_FN),
+ ("read_ranges_fn", READ_RANGES_FN),
]
diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py
index e4c0f94..c8c8ebf 100644
--- a/python/tests/test_vindex.py
+++ b/python/tests/test_vindex.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
+import ctypes
import io
import numpy as np
@@ -55,6 +56,34 @@ def reader_from_bytes(data):
return VectorIndexReader(VectorIndexInput(data))
+def test_python_read_callback_forwards_ranges_in_one_batch():
+ from paimon_vindex import _make_read_ranges_callback
+ from paimon_vindex import _ffi
+
+ class RecordingInput(VectorIndexInput):
+ def __init__(self, data):
+ super().__init__(data)
+ self.calls = []
+
+ def pread_many(self, ranges):
+ self.calls.append(list(ranges))
+ return super().pread_many(ranges)
+
+ source = RecordingInput(bytes(range(32)))
+ callback = _make_read_ranges_callback(source)
+ first = (ctypes.c_uint8 * 3)()
+ second = (ctypes.c_uint8 * 4)()
+ requests = (_ffi.PaimonVindexReadRequest * 2)(
+ _ffi.PaimonVindexReadRequest(2, first, 3),
+ _ffi.PaimonVindexReadRequest(11, second, 4),
+ )
+
+ assert callback(None, requests, 2) == 0
+ assert source.calls == [[(2, 3), (11, 4)]]
+ assert bytes(first) == bytes([2, 3, 4])
+ assert bytes(second) == bytes([11, 12, 13, 14])
+
+
def test_python_ffi_roundtrips_supported_indexes():
configs = [
(
diff --git a/tools/README.md b/tools/README.md
index e6a55a7..ff94ed5 100644
--- a/tools/README.md
+++ b/tools/README.md
@@ -1,3 +1,22 @@
+<!--
+ ~ 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.
+ -->
+
# Release tools
This directory contains helper scripts used by release managers and committers.
diff --git a/tools/check_license_headers.py b/tools/check_license_headers.py
index ec763f4..7add236 100755
--- a/tools/check_license_headers.py
+++ b/tools/check_license_headers.py
@@ -38,6 +38,7 @@ EXEMPT_FILES = {
"core/tests/fixtures/ivf_hnsw_sq_v1.hex",
"core/tests/fixtures/ivf_pq_4bit_v1.hex",
"core/tests/fixtures/ivf_pq_v1.hex",
+ "core/tests/fixtures/ivf_rq_v1.hex",
}
