xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: buffer storage write (#414)

This commit is contained in:
ZeFeng Yin 2026-06-01 10:40:49 +08:00 committed by GitHub
parent de8fb760ef
commit 74beb2a828
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 4030 additions and 731 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

452
src/ailego/buffer/vector_page_table.cc
View File

@ -13,15 +13,13 @@
// limitations under the License.
#include <algorithm>
#include <chrono>
#include <cstring>
#include <thread>
#include <ailego/utility/memory_helper.h>
#include <zvec/ailego/buffer/vector_page_table.h>
#include <zvec/core/framework/index_logger.h>
#if !defined(_MSC_VER)
#include <unistd.h>
#endif
#if defined(_MSC_VER)
#ifndef NOMINMAX
#define NOMINMAX
@ -39,6 +37,29 @@ static ssize_t zvec_pread(int fd, void *buf, size_t count, size_t offset) {
}
return static_cast<ssize_t>(bytes_read);
}
static ssize_t zvec_pwrite(int fd, const void *buf, size_t count,
size_t offset) {
HANDLE handle = reinterpret_cast<HANDLE>(_get_osfhandle(fd));
if (handle == INVALID_HANDLE_VALUE) return -1;
OVERLAPPED ov = {};
ov.Offset = static_cast<DWORD>(offset & 0xFFFFFFFF);
ov.OffsetHigh = static_cast<DWORD>(offset >> 32);
DWORD bytes_written = 0;
if (!WriteFile(handle, buf, static_cast<DWORD>(count), &bytes_written, &ov)) {
return -1;
}
return static_cast<ssize_t>(bytes_written);
}
#else
#include <unistd.h>
static inline ssize_t zvec_pread(int fd, void *buf, size_t count,
size_t offset) {
return ::pread(fd, buf, count, static_cast<off_t>(offset));
}
static inline ssize_t zvec_pwrite(int fd, const void *buf, size_t count,
size_t offset) {
return ::pwrite(fd, buf, count, static_cast<off_t>(offset));
}
#endif
namespace zvec {
@ -46,104 +67,220 @@ namespace ailego {
const size_t kVectorPageSize = MemoryHelper::PageSize();
void VectorPageTable::init(size_t entry_num) {
if (entries_) {
delete[] entries_;
bool VectorPageTable::init(size_t entry_num) {
size_t need_segments = (entry_num + kSegmentSize - 1) / kSegmentSize;
if (need_segments > kMaxSegments) {
LOG_ERROR(
"VectorPageTable::init: entry_num=%zu exceeds capacity "
"(kMaxEntries=%zu, need_segments=%zu, kMaxSegments=%zu); "
"refusing to init.",
entry_num, kMaxEntries, need_segments, kMaxSegments);
return false;
}
entry_num_ = entry_num;
entries_ = new Entry[entry_num_];
for (size_t i = 0; i < entry_num_; i++) {
entries_[i].ref_count.store(std::numeric_limits<int>::min());
entries_[i].in_evict_queue.store(false);
entries_[i].buffer = nullptr;
// Free old segments if any. init() is only called from VecBufferPool::init
// which is single-threaded with respect to other accesses, so a relaxed
// load of segment_count_ is sufficient here.
size_t old_count = segment_count_.load(std::memory_order_relaxed);
for (size_t i = 0; i < old_count; ++i) {
delete[] segments_[i];
segments_[i] = nullptr;
}
for (size_t s = 0; s < need_segments; ++s) {
segments_[s] = new Entry[kSegmentSize];
for (size_t i = 0; i < kSegmentSize; ++i) {
segments_[s][i].ref_count.store(std::numeric_limits<int>::min());
segments_[s][i].in_evict_queue.store(false);
segments_[s][i].is_dirty.store(false);
segments_[s][i].buffer = nullptr;
segments_[s][i].file_offset = 0;
}
}
// Publish new segments to readers. segment_count_ is published first
// (release) so that a reader that acquire-loads segment_count_ before
// entry_num_ also sees a consistent segment table; entry_num_ is the
// primary synchronization point used by callers via entry_num().
segment_count_.store(need_segments, std::memory_order_release);
entry_num_.store(entry_num, std::memory_order_release);
return true;
}
bool VectorPageTable::extend(size_t new_entry_num) {
// Relaxed read is fine: extend() is serialized by the caller (extend_file
// is invoked under the BufferStorage write latch). No other writer races
// with us on entry_num_ / segment_count_.
if (new_entry_num <= entry_num_.load(std::memory_order_relaxed)) {
return true;
}
size_t new_segment_count = (new_entry_num + kSegmentSize - 1) / kSegmentSize;
if (new_segment_count > kMaxSegments) {
LOG_ERROR(
"VectorPageTable::extend: new_entry_num=%zu exceeds capacity "
"(kMaxEntries=%zu, new_segment_count=%zu, kMaxSegments=%zu); "
"refusing to extend.",
new_entry_num, kMaxEntries, new_segment_count, kMaxSegments);
return false;
}
size_t old_count = segment_count_.load(std::memory_order_relaxed);
for (size_t s = old_count; s < new_segment_count; ++s) {
segments_[s] = new Entry[kSegmentSize];
for (size_t i = 0; i < kSegmentSize; ++i) {
segments_[s][i].ref_count.store(std::numeric_limits<int>::min());
segments_[s][i].in_evict_queue.store(false);
segments_[s][i].is_dirty.store(false);
segments_[s][i].buffer = nullptr;
segments_[s][i].file_offset = 0;
}
}
// Publish in the same order as init(): segment_count_ first, entry_num_
// last. Both are release-stores so that the prior segment allocation /
// Entry initialization is visible to any reader that acquire-loads either
// counter (typically via entry_num()).
segment_count_.store(new_segment_count, std::memory_order_release);
entry_num_.store(new_entry_num, std::memory_order_release);
return true;
}
char *VectorPageTable::acquire_block(block_id_t block_id) {
assert(block_id < entry_num_);
Entry &entry = entries_[block_id];
assert(block_id < entry_num_.load(std::memory_order_relaxed));
Entry &e = entry_at(block_id);
while (true) {
int current_count = entry.ref_count.load(std::memory_order_acquire);
int current_count = e.ref_count.load(std::memory_order_acquire);
if (current_count < 0) {
return nullptr;
}
if (entry.ref_count.compare_exchange_weak(current_count, current_count + 1,
std::memory_order_acq_rel,
std::memory_order_acquire)) {
return entry.buffer;
if (e.ref_count.compare_exchange_weak(current_count, current_count + 1,
std::memory_order_acq_rel,
std::memory_order_acquire)) {
return e.buffer;
}
}
}
void VectorPageTable::release_block(block_id_t block_id) {
assert(block_id < entry_num_);
Entry &entry = entries_[block_id];
assert(block_id < entry_num_.load(std::memory_order_relaxed));
Entry &e = entry_at(block_id);
if (entry.ref_count.fetch_sub(1, std::memory_order_release) == 1) {
if (e.ref_count.fetch_sub(1, std::memory_order_release) == 1) {
std::atomic_thread_fence(std::memory_order_acquire);
// Attempt to transition in_evict_queue from false -> true. The CAS ensures
// only one thread enqueues this block even if multiple threads race here.
bool expected = false;
if (entry.in_evict_queue.compare_exchange_strong(
expected, true, std::memory_order_acq_rel,
std::memory_order_relaxed)) {
if (e.in_evict_queue.compare_exchange_strong(expected, true,
std::memory_order_acq_rel,
std::memory_order_relaxed)) {
BlockEvictionQueue::BlockType block;
block.page_table = this;
block.vector_block.first = block_id;
block.vector_block.second = 0;
BlockEvictionQueue::get_instance().add_single_block(block, 0);
}
// else: block is already in the eviction queue; do not add a duplicate
// entry.
}
}
void VectorPageTable::evict_block(block_id_t block_id) {
assert(block_id < entry_num_);
Entry &entry = entries_[block_id];
char *buffer = entry.buffer;
assert(block_id < entry_num_.load(std::memory_order_relaxed));
Entry &e = entry_at(block_id);
int expected = 0;
if (entry.ref_count.compare_exchange_strong(
expected, std::numeric_limits<int>::min())) {
// Two-phase eviction to prevent data race on e.buffer with
// set_block_acquired. We first CAS to kEvicting (-1), which causes
// set_block_acquired to spin-wait; then do the actual work (flush, free,
// null buffer); finally store INT_MIN ("evicted") which unblocks
// set_block_acquired.
static constexpr int kEvicting = -1;
if (e.ref_count.compare_exchange_strong(expected, kEvicting)) {
char *buffer = e.buffer;
if (buffer && e.is_dirty.load(std::memory_order_relaxed) &&
flush_callback_) {
flush_callback_(block_id, buffer, kVectorPageSize, e.file_offset);
e.is_dirty.store(false, std::memory_order_relaxed);
}
if (buffer) {
e.buffer = nullptr;
MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize);
}
// Transition to fully-evicted state. Use release so that the
// set_block_acquired acquire-load sees e.buffer == nullptr.
e.ref_count.store(std::numeric_limits<int>::min(),
std::memory_order_release);
}
// Always reset in_evict_queue regardless of whether the CAS succeeded:
// - On success: the block is evicted; future releases should re-register it.
// - On failure: the block was re-acquired by another thread between the
// ref-count check and this call. Clearing in_evict_queue lets the next
// release_block() re-enqueue it so it is not silently lost.
entry.in_evict_queue.store(false, std::memory_order_relaxed);
e.in_evict_queue.store(false, std::memory_order_relaxed);
}
char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer) {
assert(block_id < entry_num_);
Entry &entry = entries_[block_id];
char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer,
size_t file_offset) {
assert(block_id < entry_num_.load(std::memory_order_acquire));
Entry &e = entry_at(block_id);
// Diagnostics for the kEvicting wait. The wait itself never gives up:
// the only thread that can transition kEvicting -> INT_MIN is the
// evict_block() owner, so abandoning the spin here would orphan the
// entry in kEvicting forever. Instead, we use bounded backoff and emit
// tiered logs so a stuck eviction is observable.
using clock = std::chrono::steady_clock;
const auto wait_start = clock::now();
auto last_log = wait_start;
unsigned spin_count = 0;
bool warned = false;
while (true) {
int current_count = entry.ref_count.load(std::memory_order_relaxed);
int current_count = e.ref_count.load(std::memory_order_acquire);
if (current_count >= 0) {
if (entry.ref_count.compare_exchange_weak(
current_count, current_count + 1, std::memory_order_acq_rel,
std::memory_order_acquire)) {
if (e.ref_count.compare_exchange_weak(current_count, current_count + 1,
std::memory_order_acq_rel,
std::memory_order_acquire)) {
MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize);
return entry.buffer;
return e.buffer;
}
} else if (current_count == std::numeric_limits<int>::min()) {
// Fully evicted — safe to claim this entry for our new buffer.
e.buffer = buffer;
e.file_offset = file_offset;
e.in_evict_queue.store(false, std::memory_order_relaxed);
e.is_dirty.store(false, std::memory_order_relaxed);
e.ref_count.store(1, std::memory_order_release);
return e.buffer;
} else {
entry.buffer = buffer;
entry.in_evict_queue.store(false, std::memory_order_relaxed);
entry.ref_count.store(1, std::memory_order_release);
return entry.buffer;
// kEvicting (-1): eviction is in progress on this entry.
// Tiered backoff: hot spin first, then short sleep, then longer sleep.
++spin_count;
if (spin_count < 64) {
// Pure busy wait for the common ~μs case.
} else if (spin_count < 1024) {
std::this_thread::yield();
} else if (spin_count < 8192) {
std::this_thread::sleep_for(std::chrono::microseconds(100));
} else {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
// Tiered diagnostics: warn once after 100ms, error every 1s after 1s.
const auto now = clock::now();
const auto elapsed = now - wait_start;
if (!warned && elapsed >= std::chrono::milliseconds(100)) {
LOG_WARN(
"set_block_acquired: long kEvicting wait on block_id=%zu "
"(>=100ms); evict_block may be slow",
static_cast<size_t>(block_id));
warned = true;
}
if (elapsed >= std::chrono::seconds(1) &&
(now - last_log) >= std::chrono::seconds(1)) {
const auto secs =
std::chrono::duration_cast<std::chrono::seconds>(elapsed).count();
LOG_ERROR(
"set_block_acquired: stuck in kEvicting on block_id=%zu for "
"%lld s; evict_block owner may be hung or starved",
static_cast<size_t>(block_id), static_cast<long long>(secs));
last_log = now;
}
}
}
}
VecBufferPool::VecBufferPool(const std::string &filename) {
VecBufferPool::VecBufferPool(const std::string &filename, bool writable) {
file_name_ = filename;
writable_ = writable;
#if defined(_MSC_VER)
fd_ = _open(filename.c_str(), O_RDONLY | _O_BINARY);
int flags = writable_ ? (O_RDWR | _O_BINARY) : (O_RDONLY | _O_BINARY);
fd_ = _open(filename.c_str(), flags, 0644);
#else
fd_ = open(filename.c_str(), O_RDONLY);
int flags = writable_ ? O_RDWR : O_RDONLY;
fd_ = ::open(filename.c_str(), flags, 0644);
#endif
if (fd_ < 0) {
throw std::runtime_error("Failed to open file: " + filename);
@ -164,11 +301,40 @@ VecBufferPool::VecBufferPool(const std::string &filename) {
int VecBufferPool::init() {
size_t block_num = (file_size_ + kVectorPageSize - 1) / kVectorPageSize;
page_table_.init(block_num);
if (!page_table_.init(block_num)) {
LOG_ERROR(
"VecBufferPool::init: page_table_ init failed for file[%s], "
"file_size=%zu, block_num=%zu (exceeds "
"VectorPageTable::kMaxEntries=%zu)",
file_name_.c_str(), file_size_, block_num,
VectorPageTable::kMaxEntries);
return -1;
}
block_mutexes_ =
std::make_unique<std::mutex[]>(VecBufferPool::kMutexBucketCount);
LOG_DEBUG("entry num: %zu, file_size: %zu", page_table_.entry_num(),
file_size_);
// In writable mode, inject a flush callback into the page table so that
// evict_block()/flush_block()/flush_all() can pwrite dirty blocks back to
// the backing file without needing to know about fd_ directly.
if (writable_) {
int fd = fd_;
const std::string &name = file_name_;
page_table_.set_flush_callback([fd, &name](block_id_t /*block_id*/,
char *buf, size_t sz,
size_t off) -> int {
ssize_t w = zvec_pwrite(fd, buf, sz, off);
if (w != static_cast<ssize_t>(sz)) {
LOG_ERROR(
"Buffer pool flush failed: file[%s], offset[%zu], "
"expected[%zu], got[%zd]",
name.c_str(), off, sz, w);
return -1;
}
return 0;
});
}
return 0;
}
@ -213,11 +379,7 @@ char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry) {
if (expected_bytes < kVectorPageSize) {
std::memset(buffer + expected_bytes, 0, kVectorPageSize - expected_bytes);
}
#if defined(_MSC_VER)
ssize_t read_bytes = zvec_pread(fd_, buffer, expected_bytes, page_offset);
#else
ssize_t read_bytes = pread(fd_, buffer, expected_bytes, page_offset);
#endif
if (read_bytes != static_cast<ssize_t>(expected_bytes)) {
LOG_ERROR(
"Buffer pool failed to read file at offset: file[%s], page_id[%zu], "
@ -226,15 +388,11 @@ char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry) {
MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize);
return nullptr;
}
return page_table_.set_block_acquired(page_id, buffer);
return page_table_.set_block_acquired(page_id, buffer, page_offset);
}
int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) {
#if defined(_MSC_VER)
ssize_t read_bytes = zvec_pread(fd_, buffer, length, offset);
#else
ssize_t read_bytes = pread(fd_, buffer, length, offset);
#endif
if (read_bytes != static_cast<ssize_t>(length)) {
LOG_ERROR(
"Buffer pool failed to read file at offset: file[%s], offset[%zu], "
@ -245,6 +403,141 @@ int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) {
return 0;
}
int VecBufferPool::write_range(size_t file_offset, size_t length,
const char *src) {
if (!writable_) {
LOG_ERROR("write_range called on read-only pool: file[%s]",
file_name_.c_str());
return -1;
}
if (length == 0) {
return 0;
}
size_t first_page = file_offset / kVectorPageSize;
size_t last_page = (file_offset + length - 1) / kVectorPageSize;
size_t remaining = length;
size_t src_cursor = 0;
for (size_t pg = first_page; pg <= last_page; ++pg) {
// Loading the page ensures we do not clobber unrelated bytes within the
// same page when the write is not page-aligned. acquire_buffer() pre-fills
// from the backing file (or zero-pads beyond EOF).
char *page = this->acquire_buffer(pg, 50);
if (!page) {
LOG_ERROR("write_range acquire failed: file[%s], page[%zu]",
file_name_.c_str(), pg);
return -1;
}
size_t page_start = pg * kVectorPageSize;
size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0;
size_t chunk = std::min(kVectorPageSize - intra_offset, remaining);
std::memcpy(page + intra_offset, src + src_cursor, chunk);
page_table_.mark_dirty(pg);
page_table_.release_block(pg);
src_cursor += chunk;
remaining -= chunk;
}
return 0;
}
int VecBufferPool::write_meta(size_t offset, size_t length,
const char *buffer) {
if (!writable_) {
LOG_ERROR("write_meta called on read-only pool: file[%s]",
file_name_.c_str());
return -1;
}
ssize_t w = zvec_pwrite(fd_, buffer, length, offset);
if (w != static_cast<ssize_t>(length)) {
LOG_ERROR(
"Buffer pool failed to write meta: file[%s], offset[%zu], "
"length[%zu], got[%zd]",
file_name_.c_str(), offset, length, w);
return -1;
}
return 0;
}
int VecBufferPool::flush_all() {
if (!writable_) {
return 0;
}
int rc = 0;
size_t total_dirty = 0;
size_t fail_count = 0;
for (size_t i = 0; i < page_table_.entry_num(); ++i) {
if (page_table_.is_block_dirty(i)) {
++total_dirty;
int r = page_table_.flush_block(i);
if (r != 0) {
rc = r;
++fail_count;
}
}
}
if (fail_count != 0) {
// Aggregated diagnostic so that callers (notably ~VecBufferPool, which
// discards the return value) cannot silently lose dirty pages: any
// unflushed page at this point means the on-disk image is now stale.
LOG_ERROR(
"VecBufferPool::flush_all: %zu/%zu dirty page(s) failed to flush, "
"file[%s] last_rc=%d -- on-disk data may be stale.",
fail_count, total_dirty, file_name_.c_str(), rc);
}
return rc;
}
bool VecBufferPool::extend_file(size_t new_size) {
if (!writable_) {
LOG_ERROR("extend_file called on read-only pool: file[%s]",
file_name_.c_str());
return false;
}
if (new_size <= file_size_) {
return true;
}
// Pre-validate against the page table's static capacity BEFORE mutating
// any on-disk state. Otherwise a successful ftruncate followed by a
// failed page_table_.extend() would leave the file size and the page
// table out of sync (file grew, but no Entry slots cover the new range).
size_t new_entry_num = (new_size + kVectorPageSize - 1) / kVectorPageSize;
if (new_entry_num > VectorPageTable::kMaxEntries) {
LOG_ERROR(
"extend_file: requested new_size=%zu would require %zu page entries, "
"exceeding VectorPageTable::kMaxEntries=%zu (file=%s).",
new_size, new_entry_num, VectorPageTable::kMaxEntries,
file_name_.c_str());
return false;
}
#if defined(_MSC_VER)
if (_chsize_s(fd_, static_cast<int64_t>(new_size)) != 0) {
LOG_ERROR("extend_file _chsize_s failed: file[%s], new_size[%zu]",
file_name_.c_str(), new_size);
return false;
}
#else
if (::ftruncate(fd_, static_cast<off_t>(new_size)) != 0) {
LOG_ERROR("extend_file ftruncate failed: file[%s], new_size[%zu]",
file_name_.c_str(), new_size);
return false;
}
#endif
file_size_ = new_size;
// Extend the page table to cover the new file range. Existing entries
// stay at their original addresses so concurrent readers are unaffected.
// Capacity has already been validated above, so this should never fail;
// a failure here would indicate a programming error and is logged.
if (new_entry_num > page_table_.entry_num()) {
if (!page_table_.extend(new_entry_num)) {
LOG_ERROR(
"extend_file: page_table_.extend(%zu) failed unexpectedly after "
"capacity pre-check (file=%s, new_size=%zu).",
new_entry_num, file_name_.c_str(), new_size);
return false;
}
}
return true;
}
char *VecBufferPoolHandle::get_single_page(size_t file_offset, size_t len,
size_t &out_page_id) {
size_t first_page = file_offset / kVectorPageSize;
@ -252,6 +545,10 @@ char *VecBufferPoolHandle::get_single_page(size_t file_offset, size_t len,
out_page_id = first_page;
char *page = pool_.acquire_buffer(first_page, 50);
if (!page) {
LOG_ERROR(
"VecBufferPoolHandle::get_single_page: acquire_buffer failed, "
"file_offset=%zu, len=%zu, page=%zu, page_size=%zu",
file_offset, len, first_page, kVectorPageSize);
return nullptr;
}
return page + (file_offset - first_page * kVectorPageSize);
@ -269,6 +566,11 @@ bool VecBufferPoolHandle::read_range(size_t file_offset, size_t len,
for (size_t pg = first_page; pg <= last_page; ++pg) {
char *page = pool_.acquire_buffer(pg, 50);
if (!page) {
LOG_ERROR(
"VecBufferPoolHandle::read_range: acquire_buffer failed, "
"file_offset=%zu, len=%zu, page=%zu, first_page=%zu, last_page=%zu, "
"page_size=%zu",
file_offset, len, pg, first_page, last_page, kVectorPageSize);
return false;
}
size_t page_start = pg * kVectorPageSize;
@ -286,6 +588,24 @@ int VecBufferPoolHandle::get_meta(size_t offset, size_t length, char *buffer) {
return pool_.get_meta(offset, length, buffer);
}
int VecBufferPoolHandle::write_range(size_t file_offset, size_t len,
const char *src) {
return pool_.write_range(file_offset, len, src);
}
int VecBufferPoolHandle::write_meta(size_t offset, size_t length,
const char *buffer) {
return pool_.write_meta(offset, length, buffer);
}
int VecBufferPoolHandle::flush_all() {
return pool_.flush_all();
}
bool VecBufferPoolHandle::writable() const {
return pool_.writable();
}
void VecBufferPoolHandle::release_one(block_id_t block_id) {
pool_.page_table_.release_block(block_id);
}

2
src/core/algorithm/flat/flat_streamer.cc
View File

@ -34,7 +34,7 @@ FlatStreamer<BATCH_SIZE>::FlatStreamer() : entity_(stats_) {}
template <size_t BATCH_SIZE>
FlatStreamer<BATCH_SIZE>::~FlatStreamer() {
if (state_ == STATE_INITED) {
if (state_ == STATE_INITED || state_ == STATE_OPENED) {
this->cleanup();
}
}

22
src/core/algorithm/flat/flat_streamer_entity.cc
View File

@ -165,13 +165,20 @@ int FlatStreamerEntity::add(uint64_t key, const void *vec, size_t size) {
IndexStorage::MemoryBlock head_block;
this->get_head_block(head_block);
const BlockLocation *bl =
reinterpret_cast<const BlockLocation *>(head_block.data());
if (ailego_unlikely(bl == nullptr)) {
LOG_ERROR("Failed to get block loc");
return IndexError_ReadData;
BlockLocation block;
{
const BlockLocation *bl =
reinterpret_cast<const BlockLocation *>(head_block.data());
if (ailego_unlikely(bl == nullptr)) {
LOG_ERROR("Failed to get block loc");
return IndexError_ReadData;
}
block = *bl;
}
BlockLocation block = *bl;
// Release the head block reference early so that the buffer pool ref_count
// and memory budget held by it do not block subsequent acquire/evict in this
// function (alloc_block / add_to_block may compete for the same memory).
head_block.reset(nullptr);
if (!this->is_valid_block(block)) {
int ret = this->alloc_block(block, &block);
@ -922,6 +929,9 @@ int FlatStreamerEntity::add_vector_with_id(const uint32_t id, const void *query,
this->get_head_block(head_block);
BlockLocation block =
*reinterpret_cast<const BlockLocation *>(head_block.data());
// Release buffer-pool pin before any alloc_block() call that may trigger
// append_segment() and rebuild the pool (same reason as in add()).
head_block.reset(nullptr);
if (!this->is_valid_block(block)) {
int ret = this->alloc_block(block, &block);
if (ailego_unlikely(ret != 0)) {

39
src/core/algorithm/hnsw/hnsw_index_hash.h
View File

@ -41,9 +41,9 @@ class HnswIndexHashMap {
items_(reinterpret_cast<const Item *>(data)) {}
//! Return a empty loc or the key item loc
Slot(Chunk::Pointer &&chunk, IndexStorage::MemoryBlock &&mem_block)
: chunk_(std::move(chunk)), items_block_(std::move(mem_block)) {
items_ = reinterpret_cast<const Item *>(items_block_.data());
Slot(Chunk::Pointer &&chunk, std::vector<char> &&local_data)
: chunk_(std::move(chunk)), local_data_(std::move(local_data)) {
items_ = reinterpret_cast<const Item *>(local_data_.data());
}
const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const {
auto it = &items_[key & mask];
@ -73,8 +73,8 @@ class HnswIndexHashMap {
private:
Chunk::Pointer chunk_{};
const Item *items_{nullptr}; // point to chunk data
IndexStorage::MemoryBlock items_block_{};
const Item *items_{nullptr}; // point to local_data_
std::vector<char> local_data_{};
};
public:
@ -114,9 +114,9 @@ class HnswIndexHashMap {
}
int cleanup(void) {
broker_.reset();
slots_.clear();
slots_.shrink_to_fit();
broker_.reset();
mask_bits_ = 0U;
slot_items_ = 0U;
slot_loc_mask_ = 0U;
@ -141,7 +141,6 @@ class HnswIndexHashMap {
auto idx = key >> mask_bits_;
if (idx >= slots_.size()) {
if (ailego_unlikely(idx >= slots_.capacity())) {
LOG_ERROR("no space to insert");
return false;
}
for (auto i = slots_.size(); i <= idx; ++i) {
@ -152,7 +151,6 @@ class HnswIndexHashMap {
}
auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);
if (ailego_unlikely(it == nullptr)) {
LOG_ERROR("no space to insert");
return false;
}
@ -179,14 +177,10 @@ class HnswIndexHashMap {
LOG_ERROR("Chunk resize failed, size=%zu", size);
return false;
}
//! Read the whole data to memory
IndexStorage::MemoryBlock data_block;
if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {
LOG_ERROR("Chunk read failed, size=%zu", size);
return false;
}
slots_.emplace_back(std::move(chunk), std::move(data_block));
//! Use a local zero-initialized buffer; new chunks contain all zeros,
//! so no buffer-pool read is needed and no ref_count is pinned.
std::vector<char> local_buf(size, 0);
slots_.emplace_back(std::move(chunk), std::move(local_buf));
return true;
}
@ -208,13 +202,14 @@ class HnswIndexHashMap {
i, chunk->data_size(), size);
return IndexError_InvalidFormat;
}
//! Read the whole data to memory
IndexStorage::MemoryBlock data_block;
if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {
LOG_ERROR("Chunk read failed, size=%zu", size);
return false;
//! Copy chunk data into a local buffer via fetch() so that no
//! buffer-pool block is pinned for the lifetime of the Slot.
std::vector<char> local_buf(size);
if (ailego_unlikely(chunk->fetch(0U, local_buf.data(), size) != size)) {
LOG_ERROR("Chunk fetch failed, size=%zu", size);
return IndexError_InvalidFormat;
}
slots_.emplace_back(std::move(chunk), std::move(data_block));
slots_.emplace_back(std::move(chunk), std::move(local_buf));
}
return 0;
}

2
src/core/algorithm/hnsw/hnsw_streamer.cc
View File

@ -28,7 +28,7 @@ namespace core {
HnswStreamer::HnswStreamer() = default;
HnswStreamer::~HnswStreamer() {
if (state_ == STATE_INITED) {
if (state_ == STATE_INITED || state_ == STATE_OPENED) {
this->cleanup();
}
}

20
src/core/algorithm/hnsw/hnsw_streamer_entity.cc
View File

@ -37,6 +37,7 @@ int HnswStreamerEntity::init(size_t max_doc_cnt) {
std::lock_guard<std::mutex> lock(mutex_);
broker_ = std::make_shared<ChunkBroker>(stats_);
upper_neighbor_index_ = std::make_shared<NIHashMap>();
upper_neighbor_rw_mutex_ = std::make_shared<std::shared_mutex>();
keys_map_lock_ = std::make_shared<ailego::SharedMutex>();
keys_map_ = std::make_shared<HashMap<key_t, node_id_t>>();
if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) {
@ -767,9 +768,10 @@ const HnswEntity::Pointer HnswStreamerEntity::clone() const {
HnswStreamerEntity *entity = new (std::nothrow) HnswStreamerEntity(
stats_, header(), chunk_size_, node_index_mask_bits_,
upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,
upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_,
std::move(node_chunks), std::move(upper_neighbor_chunks), broker_,
node_chunk_bases_, upper_neighbor_chunk_bases_);
upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_,
keys_map_, use_key_info_map_, std::move(node_chunks),
std::move(upper_neighbor_chunks), broker_, node_chunk_bases_,
upper_neighbor_chunk_bases_);
if (ailego_unlikely(!entity)) {
LOG_ERROR("HnswStreamerEntity new failed");
}
@ -800,9 +802,9 @@ const HnswEntity::Pointer HnswMmapStreamerEntity::clone() const {
auto *entity = new (std::nothrow) HnswMmapStreamerEntity(
stats_, header(), chunk_size_, node_index_mask_bits_,
upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,
upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_,
std::move(node_chunks), std::move(upper_neighbor_chunks), broker_,
nullptr, nullptr);
upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_,
keys_map_, use_key_info_map_, std::move(node_chunks),
std::move(upper_neighbor_chunks), broker_, nullptr, nullptr);
if (ailego_unlikely(!entity)) {
LOG_ERROR("HnswMmapStreamerEntity new failed");
}
@ -833,9 +835,9 @@ const HnswEntity::Pointer HnswContiguousStreamerEntity::clone() const {
auto *entity = new (std::nothrow) HnswContiguousStreamerEntity(
stats_, header(), chunk_size_, node_index_mask_bits_,
upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,
upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_,
std::move(node_chunks), std::move(upper_neighbor_chunks), broker_,
nullptr, nullptr);
upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_,
keys_map_, use_key_info_map_, std::move(node_chunks),
std::move(upper_neighbor_chunks), broker_, nullptr, nullptr);
if (ailego_unlikely(!entity)) {
LOG_ERROR("HnswContiguousStreamerEntity new failed");
return HnswEntity::Pointer();

71
src/core/algorithm/hnsw/hnsw_streamer_entity.h
View File

@ -17,6 +17,7 @@
#include <iostream>
#include <memory>
#include <mutex>
#include <shared_mutex>
#if defined(__linux__) || defined(__APPLE__)
#include <sys/mman.h>
#endif
@ -246,19 +247,19 @@ class HnswStreamerEntity : public HnswEntity {
using NIHashMapPointer = std::shared_ptr<NIHashMap>;
//! Clone construct, used by clone method in subclasses
HnswStreamerEntity(IndexStreamer::Stats &stats, const HNSWHeader &hd,
size_t chunk_size, uint32_t node_index_mask_bits,
uint32_t upper_neighbor_mask_bits, bool filter_same_key,
bool get_vector_enabled,
const NIHashMapPointer &upper_neighbor_index,
std::shared_ptr<ailego::SharedMutex> &keys_map_lock,
const HashMapPointer<key_t, node_id_t> &keys_map,
bool use_key_info_map,
std::vector<Chunk::Pointer> &&node_chunks,
std::vector<Chunk::Pointer> &&upper_neighbor_chunks,
const ChunkBroker::Pointer &broker,
std::shared_ptr<std::vector<const uint8_t *>> node_bases,
std::shared_ptr<std::vector<const uint8_t *>> upper_bases)
HnswStreamerEntity(
IndexStreamer::Stats &stats, const HNSWHeader &hd, size_t chunk_size,
uint32_t node_index_mask_bits, uint32_t upper_neighbor_mask_bits,
bool filter_same_key, bool get_vector_enabled,
const NIHashMapPointer &upper_neighbor_index,
const std::shared_ptr<std::shared_mutex> &upper_neighbor_rw_mutex,
std::shared_ptr<ailego::SharedMutex> &keys_map_lock,
const HashMapPointer<key_t, node_id_t> &keys_map, bool use_key_info_map,
std::vector<Chunk::Pointer> &&node_chunks,
std::vector<Chunk::Pointer> &&upper_neighbor_chunks,
const ChunkBroker::Pointer &broker,
std::shared_ptr<std::vector<const uint8_t *>> node_bases,
std::shared_ptr<std::vector<const uint8_t *>> upper_bases)
: stats_(stats),
chunk_size_(chunk_size),
node_index_mask_bits_(node_index_mask_bits),
@ -269,6 +270,7 @@ class HnswStreamerEntity : public HnswEntity {
filter_same_key_(filter_same_key),
get_vector_enabled_(get_vector_enabled),
use_key_info_map_(use_key_info_map),
upper_neighbor_rw_mutex_(upper_neighbor_rw_mutex),
upper_neighbor_index_(upper_neighbor_index),
keys_map_lock_(keys_map_lock),
keys_map_(keys_map),
@ -323,6 +325,10 @@ class HnswStreamerEntity : public HnswEntity {
inline std::pair<uint32_t, uint32_t> get_upper_neighbor_chunk_loc(
level_t level, node_id_t id) const {
// Shared lock: concurrent readers are fine, but must synchronize with
// add_upper_neighbor's exclusive lock to avoid data-race on
// slots_.size() inside HnswIndexHashMap.
std::shared_lock<std::shared_mutex> lk(*upper_neighbor_rw_mutex_);
auto it = upper_neighbor_index_->find(id);
ailego_assert_abort(it != upper_neighbor_index_->end(),
"Get upper neighbor header failed");
@ -370,6 +376,10 @@ class HnswStreamerEntity : public HnswEntity {
if (level == 0) {
return 0;
}
// Exclusive lock: protects upper_neighbor_chunks_.emplace_back() and
// upper_neighbor_index_->insert() from racing with concurrent find()
// calls in get_upper_neighbor_chunk_loc().
std::unique_lock<std::shared_mutex> lk(*upper_neighbor_rw_mutex_);
Chunk::Pointer chunk;
uint64_t chunk_offset = UINT64_MAX;
size_t neighbors_size = get_total_upper_neighbors_size(level);
@ -408,14 +418,37 @@ class HnswStreamerEntity : public HnswEntity {
meta.level = level;
meta.index = (chunk_index << upper_neighbor_mask_bits_) |
(chunk_offset / upper_neighbor_size_);
size_t zero_start = chunk_offset;
chunk_offset += upper_neighbor_size_ * level;
if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {
LOG_ERROR("HashMap insert value failed");
// IMPORTANT: order matters here.
// 1) resize so the chunk's data_size covers the new region.
// 2) zero-fill the new region: storage backends like BufferStorage do
// NOT zero on resize -- only metadata is updated, and the underlying
// page may contain stale content from a previously-evicted page.
// Without this step, NeighborsHeader::neighbor_cnt is garbage and
// select_entry_point()/search_neighbors() iterate over garbage
// node_ids, eventually triggering find()'s assertion in
// get_upper_neighbor_chunk_loc().
// 3) ONLY THEN publish the entry to upper_neighbor_index_, so that any
// concurrent reader that finds this id already sees a properly
// zeroed upper-neighbor slot.
if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {
LOG_ERROR("Chunk resize to %zu failed", (size_t)chunk_offset);
return IndexError_Runtime;
}
if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {
LOG_ERROR("Chunk resize to %zu failed", (size_t)chunk_offset);
// Use std::vector instead of a VLA: VLAs are a GNU extension and may
// produce different codegen / be rejected under clang/MSVC.
std::vector<char> zeros(neighbors_size, 0);
if (ailego_unlikely(chunk->write(zero_start, zeros.data(),
neighbors_size) != neighbors_size)) {
LOG_ERROR("Chunk write zeros failed");
return IndexError_Runtime;
}
if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {
LOG_ERROR("HashMap insert value failed");
return IndexError_Runtime;
}
@ -529,6 +562,10 @@ class HnswStreamerEntity : public HnswEntity {
protected:
IndexStreamer::Stats &stats_;
std::mutex mutex_{};
//! Guards upper_neighbor_index_ and upper_neighbor_chunks_ against
//! concurrent reads (find) and writes (insert/emplace_back).
//! Shared via shared_ptr so all clones synchronize on the SAME mutex.
mutable std::shared_ptr<std::shared_mutex> upper_neighbor_rw_mutex_{};
size_t max_index_size_{0UL};
uint32_t chunk_size_{kDefaultChunkSize};
uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize};

43
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_hash.h
View File

@ -41,9 +41,9 @@ class HnswIndexHashMap {
items_(reinterpret_cast<const Item *>(data)) {}
//! Return a empty loc or the key item loc
Slot(Chunk::Pointer &&chunk, IndexStorage::MemoryBlock &&mem_block)
: chunk_(std::move(chunk)), items_block_(std::move(mem_block)) {
items_ = reinterpret_cast<const Item *>(items_block_.data());
Slot(Chunk::Pointer &&chunk, std::vector<char> &&local_data)
: chunk_(std::move(chunk)), local_data_(std::move(local_data)) {
items_ = reinterpret_cast<const Item *>(local_data_.data());
}
const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const {
auto it = &items_[key & mask];
@ -73,8 +73,8 @@ class HnswIndexHashMap {
private:
Chunk::Pointer chunk_{};
const Item *items_{nullptr}; // point to chunk data
IndexStorage::MemoryBlock items_block_{};
const Item *items_{nullptr}; // point to local_data_
std::vector<char> local_data_{};
};
public:
@ -179,14 +179,18 @@ class HnswIndexHashMap {
LOG_ERROR("Chunk resize failed, size=%zu", size);
return false;
}
//! Read the whole data to memory
IndexStorage::MemoryBlock data_block;
if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {
LOG_ERROR("Chunk read failed, size=%zu", size);
return false;
}
slots_.emplace_back(std::move(chunk), std::move(data_block));
//! Use a local zero-initialized buffer; new chunks contain all zeros,
//! so no buffer-pool read is needed and no ref_count is pinned.
//! NOTE: Previously this used `chunk->read(0U, data_block, size)` which
//! returns a view into the underlying BufferPool page. That made the
//! Slot's `items_` pointer alias buffer-pool memory shared across
//! threads, which under clang -O3 release exposed a data race on
//! Slot::find()'s probing read of `it->second` (concurrent
//! const_cast writes from insert() were not reliably visible). Using a
//! private zero-initialized vector matches the HNSW (non-RABITQ)
//! implementation and avoids this race.
std::vector<char> local_buf(size, 0);
slots_.emplace_back(std::move(chunk), std::move(local_buf));
return true;
}
@ -208,13 +212,14 @@ class HnswIndexHashMap {
i, chunk->data_size(), size);
return IndexError_InvalidFormat;
}
//! Read the whole data to memory
IndexStorage::MemoryBlock data_block;
if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {
LOG_ERROR("Chunk read failed, size=%zu", size);
return false;
//! Copy chunk data into a local buffer via fetch() so that no
//! buffer-pool block is pinned for the lifetime of the Slot.
std::vector<char> local_buf(size);
if (ailego_unlikely(chunk->fetch(0U, local_buf.data(), size) != size)) {
LOG_ERROR("Chunk fetch failed, size=%zu", size);
return IndexError_InvalidFormat;
}
slots_.emplace_back(std::move(chunk), std::move(data_block));
slots_.emplace_back(std::move(chunk), std::move(local_buf));
}
return 0;
}

2
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc
View File

@ -40,7 +40,7 @@ HnswRabitqStreamer::HnswRabitqStreamer(IndexProvider::Pointer provider,
provider_(std::move(provider)) {}
HnswRabitqStreamer::~HnswRabitqStreamer() {
if (state_ == STATE_INITED) {
if (state_ == STATE_INITED || state_ == STATE_OPENED) {
this->cleanup();
}
}

6
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.cc
View File

@ -34,6 +34,7 @@ int HnswRabitqStreamerEntity::init(size_t max_doc_cnt) {
std::lock_guard<std::mutex> lock(mutex_);
broker_ = std::make_shared<HnswRabitqChunkBroker>(stats_);
upper_neighbor_index_ = std::make_shared<NIHashMap>();
upper_neighbor_rw_mutex_ = std::make_shared<std::shared_mutex>();
keys_map_lock_ = std::make_shared<ailego::SharedMutex>();
keys_map_ = std::make_shared<HashMap<key_t, node_id_t>>();
if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) {
@ -697,8 +698,9 @@ const HnswRabitqEntity::Pointer HnswRabitqStreamerEntity::clone() const {
new (std::nothrow) HnswRabitqStreamerEntity(
stats_, header(), chunk_size_, node_index_mask_bits_,
upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,
upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_,
std::move(node_chunks), std::move(upper_neighbor_chunks), broker_);
upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_,
keys_map_, use_key_info_map_, std::move(node_chunks),
std::move(upper_neighbor_chunks), broker_);
if (ailego_unlikely(!entity)) {
LOG_ERROR("HnswRabitqStreamerEntity new failed");
}

69
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.h
View File

@ -15,6 +15,7 @@
#pragma once
#include <iostream>
#include <shared_mutex>
#include <ailego/parallel/lock.h>
#include <sparsehash/dense_hash_map>
#include <sparsehash/dense_hash_set>
@ -216,17 +217,17 @@ class HnswRabitqStreamerEntity : public HnswRabitqEntity {
using NIHashMapPointer = std::shared_ptr<NIHashMap>;
//! Private construct, only be called by clone method
HnswRabitqStreamerEntity(IndexStreamer::Stats &stats, const HNSWHeader &hd,
size_t chunk_size, uint32_t node_index_mask_bits,
uint32_t upper_neighbor_mask_bits,
bool filter_same_key, bool get_vector_enabled,
const NIHashMapPointer &upper_neighbor_index,
std::shared_ptr<ailego::SharedMutex> &keys_map_lock,
const HashMapPointer<key_t, node_id_t> &keys_map,
bool use_key_info_map,
std::vector<Chunk::Pointer> &&node_chunks,
std::vector<Chunk::Pointer> &&upper_neighbor_chunks,
const HnswRabitqChunkBroker::Pointer &broker)
HnswRabitqStreamerEntity(
IndexStreamer::Stats &stats, const HNSWHeader &hd, size_t chunk_size,
uint32_t node_index_mask_bits, uint32_t upper_neighbor_mask_bits,
bool filter_same_key, bool get_vector_enabled,
const NIHashMapPointer &upper_neighbor_index,
const std::shared_ptr<std::shared_mutex> &upper_neighbor_rw_mutex,
std::shared_ptr<ailego::SharedMutex> &keys_map_lock,
const HashMapPointer<key_t, node_id_t> &keys_map, bool use_key_info_map,
std::vector<Chunk::Pointer> &&node_chunks,
std::vector<Chunk::Pointer> &&upper_neighbor_chunks,
const HnswRabitqChunkBroker::Pointer &broker)
: stats_(stats),
chunk_size_(chunk_size),
node_index_mask_bits_(node_index_mask_bits),
@ -237,6 +238,7 @@ class HnswRabitqStreamerEntity : public HnswRabitqEntity {
filter_same_key_(filter_same_key),
get_vector_enabled_(get_vector_enabled),
use_key_info_map_(use_key_info_map),
upper_neighbor_rw_mutex_(upper_neighbor_rw_mutex),
upper_neighbor_index_(upper_neighbor_index),
keys_map_lock_(keys_map_lock),
keys_map_(keys_map),
@ -286,6 +288,11 @@ class HnswRabitqStreamerEntity : public HnswRabitqEntity {
inline std::pair<uint32_t, uint32_t> get_upper_neighbor_chunk_loc(
level_t level, node_id_t id) const {
// Shared lock: concurrent readers are fine, but must synchronize with
// add_upper_neighbor's exclusive lock to avoid data-race on
// slots_.size() inside HnswIndexHashMap (the emplace_back in alloc_slot
// is not atomic and concurrent find() may see a stale size value).
std::shared_lock<std::shared_mutex> lk(*upper_neighbor_rw_mutex_);
auto it = upper_neighbor_index_->find(id);
ailego_assert_abort(it != upper_neighbor_index_->end(),
"Get upper neighbor header failed");
@ -334,6 +341,10 @@ class HnswRabitqStreamerEntity : public HnswRabitqEntity {
if (level == 0) {
return 0;
}
// Exclusive lock: protects upper_neighbor_chunks_.emplace_back() and
// upper_neighbor_index_->insert() from racing with concurrent find()
// calls in get_upper_neighbor_chunk_loc().
std::unique_lock<std::shared_mutex> lk(*upper_neighbor_rw_mutex_);
Chunk::Pointer chunk;
uint64_t chunk_offset = -1UL;
size_t neighbors_size = get_total_upper_neighbors_size(level);
@ -373,14 +384,37 @@ class HnswRabitqStreamerEntity : public HnswRabitqEntity {
meta.level = level;
meta.index = (chunk_index << upper_neighbor_mask_bits_) |
(chunk_offset / upper_neighbor_size_);
size_t zero_start = chunk_offset;
chunk_offset += upper_neighbor_size_ * level;
if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {
LOG_ERROR("HashMap insert value failed");
// IMPORTANT: order matters here.
// 1) resize so the chunk's data_size covers the new region.
// 2) zero-fill the new region: storage backends like BufferStorage do
// NOT zero on resize -- only metadata is updated, and the underlying
// page may contain stale content from a previously-evicted page.
// Without this step, NeighborsHeader::neighbor_cnt is garbage and
// select_entry_point()/search_neighbors() iterate over garbage
// node_ids, eventually triggering find()'s assertion in
// get_upper_neighbor_chunk_loc() at line 291.
// 3) ONLY THEN publish the entry to upper_neighbor_index_, so that any
// concurrent reader that finds this id already sees a properly
// zeroed upper-neighbor slot.
if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {
LOG_ERROR("Chunk resize to %zu failed", (size_t)chunk_offset);
return IndexError_Runtime;
}
if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {
LOG_ERROR("Chunk resize to %zu failed", (size_t)chunk_offset);
// Use std::vector instead of a VLA: VLAs are a GNU extension and may
// produce different codegen / be rejected under clang/MSVC.
std::vector<char> zeros(neighbors_size, 0);
if (ailego_unlikely(chunk->write(zero_start, zeros.data(),
neighbors_size) != neighbors_size)) {
LOG_ERROR("Chunk write zeros failed");
return IndexError_Runtime;
}
if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {
LOG_ERROR("HashMap insert value failed");
return IndexError_Runtime;
}
@ -503,6 +537,11 @@ class HnswRabitqStreamerEntity : public HnswRabitqEntity {
bool get_vector_enabled_{false};
bool use_key_info_map_{true};
// Shared via shared_ptr so that all cloned entities synchronize against
// the SAME mutex instance. A plain std::shared_mutex member would be
// independent per clone and provide no real protection for the shared
// upper_neighbor_index_ hashmap.
mutable std::shared_ptr<std::shared_mutex> upper_neighbor_rw_mutex_{};
NIHashMapPointer upper_neighbor_index_{};
mutable std::shared_ptr<ailego::SharedMutex> keys_map_lock_{};

2
src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc
View File

@ -27,7 +27,7 @@ namespace core {
HnswSparseStreamer::HnswSparseStreamer() : entity_(stats_) {}
HnswSparseStreamer::~HnswSparseStreamer() {
if (state_ == STATE_INITED) {
if (state_ == STATE_INITED || state_ == STATE_OPENED) {
this->cleanup();
}
}

2
src/core/algorithm/vamana/vamana_streamer.cc
View File

@ -26,7 +26,7 @@ namespace core {
VamanaStreamer::VamanaStreamer() = default;
VamanaStreamer::~VamanaStreamer() {
if (state_ == STATE_INITED) {
if (state_ == STATE_INITED || state_ == STATE_OPENED) {
this->cleanup();
}
}

15
src/core/interface/indexes/ivf_index.cc
View File

@ -84,15 +84,22 @@ int IVFIndex::Open(const std::string &file_path,
break;
}
case StorageOptions::StorageType::kBufferPool: {
storage_ = core::IndexFactory::CreateStorage("BufferStorage");
// NOTE: IVF index is dumped via FileDumper (plain binary file), which is
// not compatible with BufferStorage's IndexFormat layout (header/footer
// chain). Until IVF gains a BufferStorage-aware dump path, fall back to
// MMapFileReadStorage so the freshly-dumped file can be reopened.
storage_ = core::IndexFactory::CreateStorage("MMapFileReadStorage");
if (storage_ == nullptr) {
LOG_ERROR("Failed to create BufferStorage");
LOG_ERROR(
"Failed to create MMapFileReadStorage (IVF buffer-pool fallback)");
return core::IndexError_Runtime;
}
int ret = storage_->init(storage_params);
if (ret != 0) {
LOG_ERROR("Failed to init BufferStorage, path: %s, err: %s",
file_path_.c_str(), core::IndexError::What(ret));
LOG_ERROR(
"Failed to init MMapFileReadStorage (IVF buffer-pool fallback), "
"path: %s, err: %s",
file_path_.c_str(), core::IndexError::What(ret));
return ret;
}
break;

1264
src/core/utility/buffer_storage.cc
View File

File diff suppressed because it is too large Load Diff

10
src/db/index/segment/segment.cc
View File

@ -526,10 +526,20 @@ Status SegmentImpl::close() {
}
}
vector_indexers_.clear();
for (const auto &[name, indexers] : quant_vector_indexers_) {
for (auto indexer : indexers) {
indexer->Close();
}
}
quant_vector_indexers_.clear();
for (auto [name, indexer] : memory_vector_indexers_) {
indexer->Close();
}
memory_vector_indexers_.clear();
for (auto [name, indexer] : quant_memory_vector_indexers_) {
indexer->Close();
}
quant_memory_vector_indexers_.clear();
return Status::OK();
}

3
src/db/index/storage/lazy_record_batch_reader.h
View File

@ -128,7 +128,8 @@ class ParquetRecordBatchReader : public arrow::RecordBatchReader {
std::vector<std::shared_ptr<arrow::Array>> chunks(col_indices_.size());
if (with_cache_) {
for (size_t col_idx = 0; col_idx < col_indices_.size(); ++col_idx) {
auto buffer_id = ailego::ParquetBufferID(file_path_, col_idx, rg_id);
auto buffer_id =
ailego::ParquetBufferID(file_path_, col_indices_[col_idx], rg_id);
auto buffer_handle =
ailego::ParquetBufferPool::get_instance().acquire_buffer(buffer_id);
std::shared_ptr<arrow::ChunkedArray> col_chunked_array =

29
src/db/index/storage/store_helper.h
View File

@ -267,12 +267,7 @@ inline arrow::Status ConvertScalarVectorToArrayByType(
return arrow::Status::Invalid(
"Cannot convert empty vector to list array");
}
auto list_type = std::dynamic_pointer_cast<arrow::ListType>(type);
if (!list_type) {
return arrow::Status::TypeError("Expected ListType for LIST scalar");
}
auto list_type = std::static_pointer_cast<arrow::ListType>(type);
std::unique_ptr<arrow::ArrayBuilder> value_builder;
ARROW_RETURN_NOT_OK(arrow::MakeBuilder(arrow::default_memory_pool(),
list_type->value_type(),
@ -287,10 +282,9 @@ inline arrow::Status ConvertScalarVectorToArrayByType(
continue;
}
auto list_scalar = std::dynamic_pointer_cast<arrow::ListScalar>(scalar);
if (!list_scalar) {
return arrow::Status::TypeError("Expected ListScalar for LIST type");
}
// Same rationale: scalar->type->id() == LIST implies the
// scalar IS a ListScalar; avoid RTTI-dependent cast.
auto list_scalar = std::static_pointer_cast<arrow::ListScalar>(scalar);
ARROW_RETURN_NOT_OK(builder.Append());
auto value_builder_ptr = builder.value_builder();
@ -371,12 +365,10 @@ inline arrow::Status AppendFieldValueToBuilder(
}
case arrow::Type::LIST: {
auto list_builder = dynamic_cast<arrow::ListBuilder *>(builder);
auto list_type =
std::dynamic_pointer_cast<arrow::ListType>(field->type());
if (!list_type) {
return arrow::Status::TypeError("Field type is not ListType");
}
// Use static_pointer_cast: the switch guarantees type == LIST;
// dynamic_pointer_cast fails on Android due to RTTI divergence
// when Arrow is linked as a static archive.
auto list_type = std::static_pointer_cast<arrow::ListType>(field->type());
auto value_type = list_type->value_type()->id();
@ -699,8 +691,9 @@ inline arrow::Status BuildArrayFromIndicesWithType(
return BuildArrayFromIndices<arrow::BinaryArray, arrow::BinaryBuilder>(
chunked_array, indices_in_table, out_array);
case arrow::Type::LIST: {
auto list_type =
std::dynamic_pointer_cast<arrow::ListType>(col_data_type);
// static_pointer_cast: switch guarantees type == LIST; avoids
// Android RTTI divergence with Arrow static archive.
auto list_type = std::static_pointer_cast<arrow::ListType>(col_data_type);
return BuildListArrayFromIndices(chunked_array, indices_in_table,
list_type, out_array);
}

158
src/include/zvec/ailego/buffer/vector_page_table.h
View File

@ -22,6 +22,7 @@
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iostream>
#include <limits>
#include <map>
@ -48,16 +49,28 @@ class VectorPageTable {
struct Entry {
std::atomic<int> ref_count;
std::atomic<bool> in_evict_queue;
std::atomic<bool> is_dirty;
char *buffer;
size_t file_offset;
};
public:
VectorPageTable() : entry_num_(0), entries_(nullptr) {
// Callback invoked by evict_block() to persist a dirty block before its
// memory is released. Signature: (block_id, buffer, size, file_offset).
using FlushCallback = std::function<int(block_id_t, char *, size_t, size_t)>;
VectorPageTable() {
BlockEvictionQueue::get_instance().set_valid(this);
}
~VectorPageTable() {
BlockEvictionQueue::get_instance().set_invalid(this);
delete[] entries_;
// Destructor runs without concurrent readers/writers (callers guarantee
// no live handles by the time the page table is destroyed), so a relaxed
// load is sufficient here.
size_t cnt = segment_count_.load(std::memory_order_relaxed);
for (size_t i = 0; i < cnt; ++i) {
delete[] segments_[i];
}
}
VectorPageTable(const VectorPageTable &) = delete;
@ -65,7 +78,17 @@ class VectorPageTable {
VectorPageTable(VectorPageTable &&) = delete;
VectorPageTable &operator=(VectorPageTable &&) = delete;
void init(size_t entry_num);
//! Initialize the page table to cover `entry_num` entries.
//! Returns false (without modifying state) if `entry_num` exceeds the
//! statically allocated segment table capacity (kMaxEntries).
bool init(size_t entry_num);
//! Extend the page table to cover at least `new_entry_num` entries.
//! Existing entries stay at their original addresses (no invalidation).
//! Safe to call while readers operate on existing pages.
//! Returns false (without modifying state) if `new_entry_num` exceeds
//! the statically allocated segment table capacity (kMaxEntries).
bool extend(size_t new_entry_num);
char *acquire_block(block_id_t block_id);
@ -73,25 +96,101 @@ class VectorPageTable {
void evict_block(block_id_t block_id);
char *set_block_acquired(block_id_t block_id, char *buffer);
char *set_block_acquired(block_id_t block_id, char *buffer,
size_t file_offset);
void set_flush_callback(FlushCallback cb) {
flush_callback_ = std::move(cb);
}
//! Mark a loaded block as dirty so that it is persisted on eviction.
void mark_dirty(block_id_t block_id) {
assert(block_id < entry_num_.load(std::memory_order_acquire));
entry_at(block_id).is_dirty.store(true, std::memory_order_relaxed);
}
bool is_block_dirty(block_id_t block_id) const {
assert(block_id < entry_num_.load(std::memory_order_acquire));
return entry_at(block_id).is_dirty.load(std::memory_order_relaxed);
}
//! Flush a single dirty block without evicting it. Caller guarantees the
//! block is currently loaded (buffer != nullptr).
int flush_block(block_id_t block_id) {
assert(block_id < entry_num_.load(std::memory_order_acquire));
Entry &e = entry_at(block_id);
char *buffer = e.buffer;
if (!buffer || !flush_callback_) {
return 0;
}
if (!e.is_dirty.load(std::memory_order_relaxed)) {
return 0;
}
int rc = flush_callback_(block_id, buffer, kVectorPageSize, e.file_offset);
if (rc == 0) {
e.is_dirty.store(false, std::memory_order_relaxed);
}
return rc;
}
//! Returns the current number of entries. Uses acquire ordering so that
//! callers iterating over [0, entry_num()) are guaranteed to see all
//! segments_[s] writes performed by a concurrent extend()/init().
size_t entry_num() const {
return entry_num_;
return entry_num_.load(std::memory_order_acquire);
}
bool is_released(block_id_t block_id) const {
assert(block_id < entry_num_);
return entries_[block_id].ref_count.load(std::memory_order_relaxed) <= 0;
assert(block_id < entry_num_.load(std::memory_order_acquire));
return entry_at(block_id).ref_count.load(std::memory_order_relaxed) <= 0;
}
inline bool is_dead_block(BlockEvictionQueue::BlockType block) const {
Entry &entry = entries_[block.vector_block.first];
return !entry.in_evict_queue.load(std::memory_order_relaxed);
const Entry &e = entry_at(block.vector_block.first);
return !e.in_evict_queue.load(std::memory_order_relaxed);
}
private:
size_t entry_num_{0};
Entry *entries_{nullptr};
// Segmented page table: entries are split across fixed-size segments so
// that extend() can grow the table without moving existing entries.
static constexpr size_t kSegmentShift = 16; // 65536 entries per segment
static constexpr size_t kSegmentSize = size_t{1} << kSegmentShift;
static constexpr size_t kSegmentMask = kSegmentSize - 1;
public:
static constexpr size_t kMaxSegments =
2048; // up to 128M entries (512GB @ 4K)
// Maximum number of entries the segment table can ever hold. Callers
// (e.g. VecBufferPool::extend_file) can use this to pre-validate a target
// file size before mutating any on-disk state.
static constexpr size_t kMaxEntries = kMaxSegments * kSegmentSize;
private:
// entry_num_ and segment_count_ are mutated by writers in init()/extend()
// and observed by readers in entry_num() and the hot-path methods. They
// are atomic to establish a release/acquire synchronization edge with the
// (non-atomic) writes to segments_[s] performed prior to the store: any
// reader that observes the new entry_num_ is guaranteed to see the
// fully-initialized Entry slots in the corresponding segment.
std::atomic<size_t> entry_num_{0};
std::atomic<size_t> segment_count_{0};
Entry *segments_[kMaxSegments]{};
// Pair with the release-store on segment_count_ in init()/extend() so
// that any reader observing the published segment table also sees the
// fully-initialized segments_[s] pointer and Entry slots. Without this
// acquire load, segments_[s] can be re-read as nullptr or a torn
// pointer on weak memory models (and even reordered on x86 under -O2).
Entry &entry_at(size_t idx) {
(void)segment_count_.load(std::memory_order_acquire);
return segments_[idx >> kSegmentShift][idx & kSegmentMask];
}
const Entry &entry_at(size_t idx) const {
(void)segment_count_.load(std::memory_order_acquire);
return segments_[idx >> kSegmentShift][idx & kSegmentMask];
}
FlushCallback flush_callback_{};
};
class VecBufferPoolHandle;
@ -102,8 +201,11 @@ class VecBufferPool {
static constexpr size_t kMutexBucketCount = 64UL * 1024UL;
VecBufferPool(const std::string &filename);
VecBufferPool(const std::string &filename, bool writable = false);
~VecBufferPool() {
// Flush any remaining dirty blocks before tearing down memory/fd so that
// writes are not silently lost. Safe to call even in read-only mode.
(void)this->flush_all();
for (size_t i = 0; i < page_table_.entry_num(); ++i) {
assert(page_table_.is_released(i));
page_table_.evict_block(i);
@ -123,6 +225,29 @@ class VecBufferPool {
int get_meta(size_t offset, size_t length, char *buffer);
//! Write a contiguous range via the page cache; marks touched pages dirty.
//! Returns 0 on success, -1 on failure (e.g. read-only pool or I/O error).
int write_range(size_t file_offset, size_t length, const char *src);
//! Write raw bytes directly via pwrite, bypassing the page cache. Used for
//! metadata regions (header/footer/segments_meta) which are only read via
//! get_meta() and never cached.
int write_meta(size_t offset, size_t length, const char *buffer);
//! Iterate all entries and persist any dirty blocks to disk. Safe to call
//! repeatedly; no-op in read-only mode.
int flush_all();
//! Extend the backing file to `new_size` bytes via ftruncate (no-op if
//! already >= new_size), refresh the cached file_size_, and extend the
//! page_table to cover the new range. Returns true on success, false on
//! a read-only pool or I/O failure.
bool extend_file(size_t new_size);
bool writable() const {
return writable_;
}
size_t file_size() const {
return file_size_;
}
@ -131,6 +256,7 @@ class VecBufferPool {
int fd_;
size_t file_size_;
std::string file_name_;
bool writable_{false};
public:
VectorPageTable page_table_;
@ -154,6 +280,14 @@ class VecBufferPoolHandle {
int get_meta(size_t offset, size_t length, char *buffer);
int write_range(size_t file_offset, size_t len, const char *src);
int write_meta(size_t offset, size_t length, const char *buffer);
int flush_all();
bool writable() const;
void release_one(block_id_t block_id);
void acquire_one(block_id_t block_id);

69
src/include/zvec/core/framework/index_storage.h
View File

@ -14,6 +14,7 @@
#pragma once
#include <cstring>
#include <zvec/ailego/buffer/vector_page_table.h>
#include <zvec/ailego/container/params.h>
#include <zvec/core/framework/index_error.h>
@ -47,23 +48,35 @@ class IndexStorage : public IndexModule {
}
MemoryBlock(void *data) : type_(MemoryBlockType::MBT_MMAP), data_(data) {}
static MemoryBlock MakeOwned(void *owned) {
//! Build an HEAP_SCRATCH MemoryBlock that owns `owned` (allocated via
//! ailego_malloc / ailego_aligned_malloc). `size` is the byte length of
//! the buffer and is required so that copy construction / copy
//! assignment can deep-copy the buffer instead of aliasing it (a shallow
//! copy would result in use-after-free once the original block is
//! destructed and frees the buffer).
static MemoryBlock MakeOwned(void *owned, size_t size) {
MemoryBlock mb;
mb.type_ = MemoryBlockType::MBT_HEAP_SCRATCH;
mb.data_ = owned;
mb.scratch_size_ = size;
return mb;
}
MemoryBlock(const MemoryBlock &rhs) {
switch (rhs.type_) {
case MemoryBlockType::MBT_MMAP:
case MemoryBlockType::MBT_HEAP_SCRATCH:
this->reset(rhs.data_);
break;
case MemoryBlockType::MBT_BUFFERPOOL:
this->reset(rhs.buffer_pool_handle_, rhs.buffer_block_id_, rhs.data_);
buffer_pool_handle_->acquire_one(buffer_block_id_);
break;
case MemoryBlockType::MBT_HEAP_SCRATCH:
// Deep copy: each owner must hold its own buffer, otherwise the
// first destructor frees the buffer and leaves the surviving
// copies dangling.
deep_copy_from(rhs);
break;
default:
break;
}
@ -83,7 +96,9 @@ class IndexStorage : public IndexModule {
case MemoryBlockType::MBT_HEAP_SCRATCH:
type_ = MemoryBlockType::MBT_HEAP_SCRATCH;
data_ = rhs.data_;
scratch_size_ = rhs.scratch_size_;
rhs.data_ = nullptr;
rhs.scratch_size_ = 0;
rhs.type_ = MemoryBlockType::MBT_UNKNOWN;
break;
default:
@ -103,7 +118,8 @@ class IndexStorage : public IndexModule {
buffer_pool_handle_->acquire_one(buffer_block_id_);
break;
case MemoryBlockType::MBT_HEAP_SCRATCH:
this->reset(rhs.data_);
release_current();
deep_copy_from(rhs);
break;
default:
break;
@ -125,10 +141,12 @@ class IndexStorage : public IndexModule {
rhs.type_ = MemoryBlockType::MBT_UNKNOWN;
break;
case MemoryBlockType::MBT_HEAP_SCRATCH:
release_owned();
release_current();
type_ = MemoryBlockType::MBT_HEAP_SCRATCH;
data_ = rhs.data_;
scratch_size_ = rhs.scratch_size_;
rhs.data_ = nullptr;
rhs.scratch_size_ = 0;
rhs.type_ = MemoryBlockType::MBT_UNKNOWN;
break;
default:
@ -154,6 +172,7 @@ class IndexStorage : public IndexModule {
break;
}
data_ = nullptr;
scratch_size_ = 0;
}
const void *data() const {
@ -188,6 +207,10 @@ class IndexStorage : public IndexModule {
void *data_{nullptr};
mutable ailego::VecBufferPoolHandle *buffer_pool_handle_{nullptr};
size_t buffer_block_id_{0};
//! Byte size of the heap-scratch buffer pointed to by `data_`; only used
//! when type_ == MBT_HEAP_SCRATCH. Required for safe deep-copy on
//! copy-construction / copy-assignment of HEAP_SCRATCH blocks.
size_t scratch_size_{0};
private:
void release_owned() {
@ -195,6 +218,44 @@ class IndexStorage : public IndexModule {
ailego_free(data_);
data_ = nullptr;
}
scratch_size_ = 0;
}
//! Drop whatever the current MemoryBlock holds, regardless of type, so
//! that the slot is ready to receive new ownership. Mirrors what the
//! destructor would do (minus zeroing data_) but leaves the type alone
//! for the caller to overwrite immediately afterwards.
void release_current() {
switch (type_) {
case MemoryBlockType::MBT_BUFFERPOOL:
if (buffer_pool_handle_) {
buffer_pool_handle_->release_one(buffer_block_id_);
buffer_pool_handle_ = nullptr;
}
break;
case MemoryBlockType::MBT_HEAP_SCRATCH:
release_owned();
break;
default:
break;
}
data_ = nullptr;
type_ = MemoryBlockType::MBT_UNKNOWN;
}
//! Allocate a fresh buffer of the same size as `rhs.scratch_size_`,
//! memcpy `rhs.data_` into it, and become the new owner. Used by the
//! HEAP_SCRATCH copy ctor / copy assignment so the original and the
//! copy each free their own buffer independently.
void deep_copy_from(const MemoryBlock &rhs) {
type_ = MemoryBlockType::MBT_HEAP_SCRATCH;
scratch_size_ = rhs.scratch_size_;
if (scratch_size_ > 0 && rhs.data_) {
data_ = ailego_malloc(scratch_size_);
std::memcpy(data_, rhs.data_, scratch_size_);
} else {
data_ = nullptr;
}
}
};

246
tests/core/algorithm/flat/flat_streamer_buffer_test.cc
View File

@ -168,6 +168,251 @@ TEST_F(FlatStreamerTest, TestLinearSearch) {
read_streamer.reset();
}
TEST_F(FlatStreamerTest, TestLinearSearchBuffer) {
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
IndexStreamer::Pointer write_streamer =
IndexFactory::CreateStreamer("FlatStreamer");
ASSERT_TRUE(write_streamer != nullptr);
Params params;
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
auto storage = IndexFactory::CreateStorage("BufferStorage");
ASSERT_NE(nullptr, storage);
Params stg_params;
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchBuffer", true));
ASSERT_EQ(0, write_streamer->open(storage));
auto ctx = write_streamer->create_context();
ASSERT_TRUE(!!ctx);
size_t cnt = 10000UL;
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
for (size_t i = 0; i < cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
}
write_streamer->flush(0UL);
write_streamer->close();
write_streamer.reset();
storage->close();
IndexStreamer::Pointer read_streamer =
IndexFactory::CreateStreamer("FlatStreamer");
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
ASSERT_NE(nullptr, read_storage);
ASSERT_EQ(0, read_storage->init(stg_params));
ASSERT_EQ(0, read_storage->open(dir_ + "Test/LinearSearchBuffer", false));
ASSERT_EQ(0, read_streamer->open(read_storage));
size_t topk = 3;
auto provider = read_streamer->create_provider();
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_FLOAT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
ctx->set_topk(100U);
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = 10.1f;
}
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
ASSERT_EQ(10, result[0].key());
ASSERT_EQ(11, result[1].key());
ASSERT_EQ(5, result[10].key());
ASSERT_EQ(0, result[20].key());
ASSERT_EQ(30, result[30].key());
ASSERT_EQ(35, result[35].key());
ASSERT_EQ(99, result[99].key());
ElapsedTime elapsed_time;
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_FLOAT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
read_streamer->close();
read_streamer.reset();
}
TEST_F(FlatStreamerTest, TestLinearSearchBufferMMap) {
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
IndexStreamer::Pointer write_streamer =
IndexFactory::CreateStreamer("FlatStreamer");
ASSERT_TRUE(write_streamer != nullptr);
Params params;
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
auto storage = IndexFactory::CreateStorage("BufferStorage");
ASSERT_NE(nullptr, storage);
Params stg_params;
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchBuffer", true));
ASSERT_EQ(0, write_streamer->open(storage));
auto ctx = write_streamer->create_context();
ASSERT_TRUE(!!ctx);
size_t cnt = 10000UL;
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
for (size_t i = 0; i < cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
}
write_streamer->flush(0UL);
write_streamer->close();
write_streamer.reset();
storage->close();
IndexStreamer::Pointer read_streamer =
IndexFactory::CreateStreamer("FlatStreamer");
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, read_storage);
ASSERT_EQ(0, read_storage->init(stg_params));
ASSERT_EQ(0, read_storage->open(dir_ + "Test/LinearSearchBuffer", false));
ASSERT_EQ(0, read_streamer->open(read_storage));
size_t topk = 3;
auto provider = read_streamer->create_provider();
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_FLOAT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
ctx->set_topk(100U);
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = 10.1f;
}
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
ASSERT_EQ(10, result[0].key());
ASSERT_EQ(11, result[1].key());
ASSERT_EQ(5, result[10].key());
ASSERT_EQ(0, result[20].key());
ASSERT_EQ(30, result[30].key());
ASSERT_EQ(35, result[35].key());
ASSERT_EQ(99, result[99].key());
ElapsedTime elapsed_time;
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_FLOAT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
read_streamer->close();
read_streamer.reset();
}
TEST_F(FlatStreamerTest, TestLinearSearchWithLRU) {
MemoryLimitPool::get_instance().init(100 * 1024UL * 1024UL);
#ifdef __ANDROID__
@ -351,7 +596,6 @@ TEST_F(FlatStreamerTest, TestLinearSearchMMap) {
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
const float *data = (float *)provider->get_vector(result1[0].key());
EXPECT_FLOAT_EQ(data[j], i);

248
tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc
View File

@ -171,6 +171,254 @@ TEST_F(HnswStreamerTest, TestHnswSearch) {
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
}
TEST_F(HnswStreamerTest, TestHnswSearchBuffer) {
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
IndexStreamer::Pointer write_streamer =
IndexFactory::CreateStreamer("HnswStreamer");
ASSERT_TRUE(write_streamer != nullptr);
Params params;
params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true);
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
auto storage = IndexFactory::CreateStorage("BufferStorage");
ASSERT_NE(nullptr, storage);
Params stg_params;
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "Test/TestHnswSearchBuffer", true));
ASSERT_EQ(0, write_streamer->open(storage));
auto ctx = write_streamer->create_context();
ASSERT_TRUE(!!ctx);
size_t cnt = 10000UL;
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
for (size_t i = 0; i < cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
}
write_streamer->flush(0UL);
write_streamer->close();
write_streamer.reset();
storage->close();
IndexStreamer::Pointer read_streamer =
IndexFactory::CreateStreamer("HnswStreamer");
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
ASSERT_NE(nullptr, read_storage);
ASSERT_EQ(0, read_storage->init(stg_params));
ASSERT_EQ(0, read_storage->open(dir_ + "Test/TestHnswSearchBuffer", false));
ASSERT_EQ(0, read_streamer->open(read_storage));
size_t topk = 3;
auto provider = read_streamer->create_provider();
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
ctx->set_topk(100U);
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = 10.1f;
}
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
ASSERT_EQ(10, result[0].key());
ASSERT_EQ(11, result[1].key());
ASSERT_EQ(5, result[10].key());
ASSERT_EQ(0, result[20].key());
ASSERT_EQ(30, result[30].key());
ASSERT_EQ(35, result[35].key());
ASSERT_EQ(99, result[99].key());
ElapsedTime elapsed_time;
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
read_streamer->close();
read_streamer.reset();
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
}
TEST_F(HnswStreamerTest, TestHnswSearchBufferMMap) {
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
IndexStreamer::Pointer write_streamer =
IndexFactory::CreateStreamer("HnswStreamer");
ASSERT_TRUE(write_streamer != nullptr);
Params params;
params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true);
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
auto storage = IndexFactory::CreateStorage("BufferStorage");
ASSERT_NE(nullptr, storage);
Params stg_params;
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "Test/TestHnswSearchBufferMMap", true));
ASSERT_EQ(0, write_streamer->open(storage));
auto ctx = write_streamer->create_context();
ASSERT_TRUE(!!ctx);
size_t cnt = 10000UL;
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
for (size_t i = 0; i < cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
}
write_streamer->flush(0UL);
write_streamer->close();
write_streamer.reset();
storage->close();
IndexStreamer::Pointer read_streamer =
IndexFactory::CreateStreamer("HnswStreamer");
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, read_storage);
ASSERT_EQ(0, read_storage->init(stg_params));
ASSERT_EQ(0, read_storage->open(dir_ + "Test/TestHnswSearchBufferMMap", false));
ASSERT_EQ(0, read_streamer->open(read_storage));
size_t topk = 3;
auto provider = read_streamer->create_provider();
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
ctx->set_topk(100U);
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = 10.1f;
}
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
ASSERT_EQ(10, result[0].key());
ASSERT_EQ(11, result[1].key());
ASSERT_EQ(5, result[10].key());
ASSERT_EQ(0, result[20].key());
ASSERT_EQ(30, result[30].key());
ASSERT_EQ(35, result[35].key());
ASSERT_EQ(99, result[99].key());
ElapsedTime elapsed_time;
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = i;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(topk, result1.size());
IndexStorage::MemoryBlock block;
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
const float *data = (float *)block.data();
for (size_t j = 0; j < dim; ++j) {
ASSERT_EQ(data[j], i);
}
ASSERT_EQ(i, result1[0].key());
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
ctx->set_topk(topk);
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(i, result2[0].key());
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
}
read_streamer->close();
read_streamer.reset();
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
}
TEST_F(HnswStreamerTest, TestHnswSearchMMap) {
IndexStreamer::Pointer write_streamer =
IndexFactory::CreateStreamer("HnswStreamer");

1181
tests/core/utility/buffer_storage_write_test.cc Normal file
View File

File diff suppressed because it is too large Load Diff

806
tests/db/collection_test.cc
View File

File diff suppressed because it is too large Load Diff