From 8ce8e3e228383203dd84d2ac1441ad009a1703db Mon Sep 17 00:00:00 2001 From: ZeFeng Yin Date: Tue, 2 Jun 2026 09:51:41 +0800 Subject: [PATCH] chore: rm buffer manager (#437) --- src/ailego/buffer/buffer_manager.cc | 603 ------------------ src/db/common/global_resource.cc | 1 - .../index/storage/bufferpool_forward_store.cc | 1 - .../index/storage/bufferpool_forward_store.h | 4 - .../index/storage/lazy_record_batch_reader.h | 1 - .../zvec/ailego/buffer/buffer_manager.h | 263 -------- .../zvec/ailego/buffer/parquet_hash_table.h | 2 +- tests/ailego/buffer/buffer_manager_test.cc | 134 ---- .../flat/flat_streamer_buffer_test.cc | 1 - .../flat/flat_streamer_buffer_time_test.cc | 1 - .../core/algorithm/flat/flat_streamer_test.cc | 1 - .../flat_sparse_streamer_buffer_test.cc | 1 - .../flat_sparse/flat_sparse_streamer_test.cc | 1 - .../hnsw/hnsw_streamer_buffer_test.cc | 1 - tests/core/interface/index_interface_test.cc | 2 - tests/core/utility/buffer_storage_test.cpp | 1 - .../core/utility/buffer_storage_write_test.cc | 1 - tests/core/utility/mmap_file_storage_test.cpp | 1 - tests/db/index/segment/segment_test.cc | 1 - 19 files changed, 1 insertion(+), 1020 deletions(-) delete mode 100644 src/ailego/buffer/buffer_manager.cc delete mode 100644 src/include/zvec/ailego/buffer/buffer_manager.h delete mode 100644 tests/ailego/buffer/buffer_manager_test.cc diff --git a/src/ailego/buffer/buffer_manager.cc b/src/ailego/buffer/buffer_manager.cc deleted file mode 100644 index be3b675..0000000 --- a/src/ailego/buffer/buffer_manager.cc +++ /dev/null @@ -1,603 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed 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. - -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __clang__ -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wunused-parameter" -#pragma clang diagnostic ignored "-Wshadow" -#elif defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-parameter" -#pragma GCC diagnostic ignored "-Wshadow" -#endif - -#include - -#ifdef __clang__ -#pragma clang diagnostic pop -#elif defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic pop -#endif - - -namespace zvec { - - -namespace ailego { - - -namespace { - - -struct IDHash { - size_t operator()(const BufferID &buffer_id) const { - size_t hash = std::hash{}(static_cast(buffer_id.type)); - hash = hash ^ (std::hash{}(buffer_id.file_id)); - if (buffer_id.type == BufferID::TYPE::PARQUET) { - hash = hash * 31 + std::hash{}(buffer_id.parquet().column); - hash = hash * 31 + std::hash{}(buffer_id.parquet().row_group); - } else if (buffer_id.type == BufferID::TYPE::VECTOR) { - hash = hash * 31 + std::hash{}(buffer_id.vector().offset); - } - return hash; - } -}; - - -struct IDEqual { - bool operator()(const BufferID &a, const BufferID &b) const { - if (a.type != b.type) { - return false; - } - if (a.file_name != b.file_name) { - return false; - } - if (a.file_id != b.file_id) { - return false; - } - if (a.mtime != b.mtime) { - return false; - } - if (a.type == BufferID::TYPE::PARQUET) { - return a.parquet().column == b.parquet().column && - a.parquet().row_group == b.parquet().row_group; - } else if (a.type == BufferID::TYPE::VECTOR) { - return a.vector().offset == b.vector().offset; - } else { - return false; - } - } -}; - - -} // namespace - - -struct BufferManager::BufferContext { - BufferContext(const BufferID &id, BufferPool *p) : id(id), pool(p) {}; - BufferContext(const BufferContext &) = delete; - BufferContext(BufferContext &&) = delete; - BufferContext &operator=(const BufferContext &) = delete; - BufferContext &operator=(BufferContext &&) = delete; - - - ~BufferContext() { - if (vector) { - ailego_aligned_free(vector); - } - } - - - typedef std::unique_ptr Pointer; - - - enum State : uint32_t { - IDLE = 0, // Empty and not held by any users, not in LRU - RESERVED = 1, // Pinned by a user but no data yet, not in LRU - IN_USE = 2, // Pinned by a user and data is present, not in LRU - CACHED = 3, // Data is present but not held by any users, in LRU - ERROR = 4 // Something went wrong, not in LRU - }; - - - // Identifier for the buffer - BufferID id; - - // Current state - State state{IDLE}; - - // The size of the buffer - uint32_t size{0}; - - // Handle of the file backing this buffer - File file; - - // The number of external references to this buffer (via pin/unpin) - std::atomic refs_buf{0}; - - // The number of external references to this context (via BufferHandle) - std::atomic refs_context{0}; - - BufferPool *pool{nullptr}; - - // A shared pointer to the buffers allocated for arrow parquet data - std::shared_ptr arrow{nullptr}; - - // Guard original arrow buffers to prevent premature deletion - std::vector> arrow_refs{}; - - // A pointer to the buffer allocated for vector data - void *vector{nullptr}; - - // Doubly linked LRU list - BufferContext *next{nullptr}; - BufferContext *prev{nullptr}; - - - // Return a string representation of the status - const std::string status_string() const; - - // Populate the buffer with parquet data - arrow::Status read_arrow_parquet(); - - // Populate the buffer with vector data - bool read_vector(); -}; - - -const std::string BufferManager::BufferContext::status_string() const { - std::string msg{id.to_string() + ": "}; - switch (state) { - case State::IDLE: { - msg += "Idle"; - break; - } - case State::RESERVED: { - msg += "Reserved"; - break; - } - case State::IN_USE: { - msg += "In use"; - break; - } - case State::CACHED: { - msg += "Cached"; - break; - } - case State::ERROR: { - msg += "Error"; - break; - } - } - return msg; -} - - -arrow::Status BufferManager::BufferContext::read_arrow_parquet() { - // TODO: file handler and memory pool can be optimized - arrow::MemoryPool *mem_pool = arrow::default_memory_pool(); - - // Open file - std::shared_ptr input; - const auto &file_name = id.file_name; - ARROW_ASSIGN_OR_RAISE(input, arrow::io::ReadableFile::Open(file_name)); - - // Open reader - std::unique_ptr reader; - ARROW_ASSIGN_OR_RAISE(reader, parquet::arrow::OpenFile(input, mem_pool)); - - // Perform read - int row_group = id.parquet().row_group; - int column = id.parquet().column; - auto s = reader->RowGroup(row_group)->Column(column)->Read(&arrow); - if (!s.ok()) { - LOG_ERROR("Failed to read parquet file[%s]", file_name.c_str()); - arrow = nullptr; - return s; - } - - // Compute the memory usage and hijack Arrow's buffers with our implementation - for (auto &array : arrow->chunks()) { - auto &buffers = array->data()->buffers; - for (size_t buf_idx = 0; buf_idx < buffers.size(); ++buf_idx) { - if (buffers[buf_idx] == nullptr) { - continue; - } - // Keep references to original buffers to prevent premature deletion - arrow_refs.emplace_back(buffers[buf_idx]); - size += buffers[buf_idx]->capacity(); - // Create hijacked buffer with custom deleter that notifies us when Arrow - // is finished with the buffer - std::shared_ptr hijacked_buffer( - buffers[buf_idx].get(), BufferManager::ArrowBufferDeleter(this)); - buffers[buf_idx] = hijacked_buffer; - } - } - - return arrow::Status::OK(); -} - - -bool BufferManager::BufferContext::read_vector() { - const auto &file_name = id.file_name; - if (!file.is_valid()) { - if (!File::IsExist(file_name)) { - LOG_ERROR("File[%s] does not exist", file_name.c_str()); - return false; - } - if (!File::IsRegular(file_name)) { - LOG_ERROR("[%s] is not a regular file", file_name.c_str()); - return false; - } - if (!file.open(file_name.c_str(), true, false)) { - LOG_ERROR("Failed to open file[%s]", file_name.c_str()); - return false; - } - } - AILEGO_DEFER([this] { file.close(); }); - uint32_t len = id.vector().length; - vector = (uint8_t *)ailego_aligned_malloc(len, 64); // 64-byte alignment - if (vector == nullptr) { - LOG_ERROR("Failed to allocate buffer for file[%s]", file_name.c_str()); - return false; - } - uint32_t offset = id.vector().offset; - if (file.read(offset, vector, len) != len) { - LOG_ERROR("Failed to read file[%s]", file_name.c_str()); - ailego_aligned_free(vector); - vector = nullptr; - return false; - } - size = len; - return true; -} - - -// Thread-safe buffer pool implementation. -// -// BufferContext states: -// 1. Must exist in the lookup (hash) table. -// 2. LRU list presence: -// - In LRU: holds memory but not pinned by any users -// - Not in LRU: either holds memory pinned by users, or doesn't hold memory -// 3. External references: when an external user acquires a context and pins the -// memory, that context is removed from LRU list; when they unpins the -// memory, that context is moved to LRU list if it was the last reference. -// -// Any operation on the hash table is protected by mutex_table_. -// Any change to context state and LRU list is protected by mutex_context_. -// -class BufferManager::BufferPool { - public: - explicit BufferPool(uint64_t limit) : limit_(limit) { - sentinel_.next = &sentinel_; - sentinel_.prev = &sentinel_; - } - - - BufferContext *acquire_locked(BufferID &id) { - std::lock_guard lock(mutex_context_); - if (auto iter = table_.find(id); iter != table_.end()) { - return iter->second.get(); - } - auto [iter, _] = - table_.emplace(id, std::make_unique(id, this)); - return iter->second.get(); - } - - - void try_release_context_locked(BufferContext *context) { - if (context->refs_context.load() != 0) { - return; - } - std::lock_guard lock(mutex_table_); - if (context->refs_context.load() != 0) { - return; - } - if (context->state == BufferContext::State::IDLE) { - table_.erase(context->id); - } - } - - - void pin_locked(BufferContext *ctx) { - std::lock_guard lock(mutex_context_); - if (ctx->state == BufferContext::State::IDLE) { - return pin_at_IDLE(ctx); - } - if (ctx->state == BufferContext::State::IN_USE) { - return pin_at_IN_USE(ctx); - } - if (ctx->state == BufferContext::State::CACHED) { - return pin_at_CACHED(ctx); - } - if (ctx->state == BufferContext::State::ERROR) { - return; - } - } - - - bool unpin_locked(BufferContext *ctx) { - uint32_t prev_refs = ctx->refs_buf.fetch_sub(1); - if (prev_refs > 1) { - return false; - } - std::lock_guard lock(mutex_context_); - if (ctx->refs_buf.load() == 0 && - ctx->state != BufferContext::State::CACHED) { - ctx->state = BufferContext::State::CACHED; - LRU_insert(ctx); - return true; - } else { - return false; - } - } - - - void LRU_insert_locked(BufferContext *context) { - std::lock_guard lock(mutex_context_); - LRU_insert(context); - } - - - void LRU_remove_locked(BufferContext *context) { - std::lock_guard lock(mutex_context_); - LRU_remove(context); - } - - - uint64_t usage() const { - return usage_; - } - - - private: - void pin_at_IDLE(BufferContext *ctx) { - ctx->state = BufferContext::State::RESERVED; - - while (usage_ >= limit_) { - // The tail of LRU list is the least recently used context - BufferContext *victim = sentinel_.prev; - if (victim == &sentinel_) { // No victim could be found - ctx->state = BufferContext::State::ERROR; - return; - } - if (victim->state == BufferContext::State::ERROR) { - LRU_remove(victim); - try_release_context_locked(ctx); - continue; - } - if (victim->id.type == BufferID::TYPE::PARQUET) { - victim->arrow_refs.clear(); - } else { - ailego_aligned_free(victim->vector); - victim->vector = nullptr; - } - victim->state = BufferContext::State::IDLE; - LRU_remove(victim); - try_release_context_locked(ctx); - usage_ -= victim->size; - } - - if (ctx->id.type == BufferID::TYPE::PARQUET) { - if (ctx->read_arrow_parquet().ok()) { - ctx->state = BufferContext::State::IN_USE; - ctx->refs_buf.fetch_add(ctx->arrow_refs.size()); - usage_ += ctx->size; - } else { - LOG_ERROR("Failed to read to %s", ctx->id.to_string().c_str()); - ctx->state = BufferContext::State::ERROR; - } - } else { - if (ctx->read_vector()) { - ctx->state = BufferContext::State::IN_USE; - ctx->refs_buf.fetch_add(1); - usage_ += ctx->size; - } else { - LOG_ERROR("Failed to read to %s", ctx->id.to_string().c_str()); - ctx->state = BufferContext::State::ERROR; - } - } - } - - - void pin_at_IN_USE(BufferContext *ctx) { - if (ctx->id.type == BufferID::TYPE::PARQUET) { - ctx->refs_buf.fetch_add(ctx->arrow_refs.size()); - } else { - ctx->refs_buf.fetch_add(1); - } - } - - - void pin_at_CACHED(BufferContext *ctx) { - if (ctx->id.type == BufferID::TYPE::PARQUET) { - ctx->refs_buf.fetch_add(ctx->arrow_refs.size()); - } else { - ctx->refs_buf.fetch_add(1); - } - LRU_remove(ctx); - ctx->state = BufferContext::State::IN_USE; - } - - - void LRU_insert(BufferContext *context) { - if (context->refs_buf > 0) { - return; // Already pinned, should not be evicted - } - if (context->next != nullptr || context->prev != nullptr) { - return; - } - // Insert the context to the head of LRU list - context->next = sentinel_.next; - context->prev = &sentinel_; - sentinel_.next = context; - context->next->prev = context; - inactive_ += context->size; - } - - - void LRU_remove(BufferContext *context) { - if (context->next == nullptr) { - return; // Not in LRU list - } - context->next->prev = context->prev; - context->prev->next = context->next; - context->next = nullptr; - context->prev = nullptr; - inactive_ -= context->size; - } - - private: - using Table = - std::unordered_map; - - uint64_t limit_; - std::atomic usage_{0}; - std::atomic inactive_{0}; - - Table table_{}; - std::mutex mutex_table_{}; - BufferContext sentinel_{BufferID{}, this}; // LRU list sentinel - std::mutex mutex_context_{}; -}; - - -BufferManager::ArrowBufferDeleter::ArrowBufferDeleter(BufferContext *c) - : context(c) {} - - -void BufferManager::ArrowBufferDeleter::operator()(arrow::Buffer *) { - context->pool->unpin_locked(context); -} - - -BufferHandle::BufferHandle(BufferContext *context) : context_(context) { - if (context_ != nullptr) { - pool_ = context_->pool; - context_->refs_context.fetch_add(1); - } -} - - -BufferHandle::~BufferHandle() { - if (context_ != nullptr) { - uint32_t prev_refs = context_->refs_context.fetch_sub(1); - if (prev_refs > 1) { - return; - } - if (context_->state == BufferContext::State::IDLE) { - pool_->try_release_context_locked(context_); - } - } -} - - -std::shared_ptr BufferHandle::pin_parquet_data() { - pool_->pin_locked(context_); - return context_->arrow; -} - - -void *BufferHandle::pin_vector_data() { - if (!context_) { - return nullptr; - } - pool_->pin_locked(context_); - return context_->vector; -} - - -bool BufferHandle::unpin_vector_data() { - if (!context_) { - return true; - } - return pool_->unpin_locked(context_); -} - - -uint32_t BufferHandle::references() const { - return context_->refs_buf.load(); -} - - -uint32_t BufferHandle::size() const { - return context_->size; -} - - -void BufferManager::init(uint64_t limit, uint32_t num_shards) { - pools_.clear(); - uint64_t limit_per_shard = ailego_align(limit / num_shards, 4096); - for (uint32_t i = 0; i < num_shards; ++i) { - auto pool = new BufferPool(limit_per_shard); - pools_.push_back(pool); - } - LOG_INFO( - "BufferManager initialized with [%u] buffer pools, [%zu] bytes memory " - "limit per pool, total memory limit [%zu] bytes", - num_shards, (size_t)limit_per_shard, (size_t)limit); -} - - -BufferHandle BufferManager::acquire(BufferID &buffer_id) { - static IDHash id_hash{}; - auto hash_val = id_hash(buffer_id); - auto ctx = pools_[hash_val % pools_.size()]->acquire_locked(buffer_id); - return BufferHandle(ctx); -} - - -std::unique_ptr BufferManager::acquire_ptr(BufferID &buffer_id) { - static IDHash id_hash{}; - auto hash_val = id_hash(buffer_id); - auto ctx = pools_[hash_val % pools_.size()]->acquire_locked(buffer_id); - return std::make_unique(ctx); -} - - -uint64_t BufferManager::total_size_in_bytes() const { - uint64_t total_usage = 0; - for (auto pool : pools_) { - total_usage += pool->usage(); - } - return total_usage; -} - - -void BufferManager::cleanup() { - for (auto pool : pools_) { - delete pool; - } - pools_.clear(); -} - -BufferManager::~BufferManager() { - cleanup(); -} - - -} // namespace ailego - - -} // namespace zvec \ No newline at end of file diff --git a/src/db/common/global_resource.cc b/src/db/common/global_resource.cc index 6c079d4..d6c7a4b 100644 --- a/src/db/common/global_resource.cc +++ b/src/db/common/global_resource.cc @@ -14,7 +14,6 @@ #include "db/common/global_resource.h" #include #include -#include #include namespace zvec { diff --git a/src/db/index/storage/bufferpool_forward_store.cc b/src/db/index/storage/bufferpool_forward_store.cc index f777507..36aefe4 100644 --- a/src/db/index/storage/bufferpool_forward_store.cc +++ b/src/db/index/storage/bufferpool_forward_store.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include "db/index/storage/store_helper.h" diff --git a/src/db/index/storage/bufferpool_forward_store.h b/src/db/index/storage/bufferpool_forward_store.h index c841055..45ddc3a 100644 --- a/src/db/index/storage/bufferpool_forward_store.h +++ b/src/db/index/storage/bufferpool_forward_store.h @@ -23,7 +23,6 @@ #include #include #include -#include #include #include "base_forward_store.h" @@ -123,9 +122,6 @@ class BufferPoolForwardStore /// Number of rows in each row group std::vector row_group_row_nums_; - - /// Buffer manager for caching data - std::shared_ptr buffer_manager_; }; } // namespace zvec \ No newline at end of file diff --git a/src/db/index/storage/lazy_record_batch_reader.h b/src/db/index/storage/lazy_record_batch_reader.h index e1286e3..8102995 100644 --- a/src/db/index/storage/lazy_record_batch_reader.h +++ b/src/db/index/storage/lazy_record_batch_reader.h @@ -16,7 +16,6 @@ #include #include -#include #include #include "db/common/constants.h" diff --git a/src/include/zvec/ailego/buffer/buffer_manager.h b/src/include/zvec/ailego/buffer/buffer_manager.h deleted file mode 100644 index 1e34961..0000000 --- a/src/include/zvec/ailego/buffer/buffer_manager.h +++ /dev/null @@ -1,263 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed 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. - - -#pragma once - - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace arrow { -class ChunkedArray; -class Array; -class DataType; -class Scalar; -template -class Result; -class Status; -class Buffer; -} // namespace arrow - -namespace zvec { - - -namespace ailego { - - -struct BufferID; -class BufferManager; -class BufferHandle; - - -struct BufferID { - struct ParquetPos { - int column; - int row_group; - }; - struct VectorPos { - uint32_t offset; - uint32_t length; - }; - union Position { - explicit Position() = default; - ParquetPos forward; - VectorPos vector; - }; - enum TYPE { - PARQUET = 1, - VECTOR = 2, - UNKNOWN = 0, - }; - - - static std::uint64_t getLastModifiedNs(const std::filesystem::path &p) { - auto ftime = std::filesystem::last_write_time(p); - return static_cast(ftime.time_since_epoch().count()); - } - - // Cross-platform helper to get nanosecond modification time - // static long get_st_mtime_nsec(const struct stat &file_stat) { - // #ifdef __APPLE__ - // return file_stat.st_mtim.tv_nsec; - // #else - // return file_stat.st_mtim.tv_nsec; - // #endif - // } - - static BufferID ParquetID(const std::string &file_name, int column, - int row_group) { - BufferID buffer_id{}; - buffer_id.type = TYPE::PARQUET; - buffer_id.file_name = file_name; - buffer_id.pos.forward.column = column; - buffer_id.pos.forward.row_group = row_group; - struct stat file_stat; - if (stat(file_name.c_str(), &file_stat) == 0) { - // file_stat.st_ino contains the inode number - // file_stat.st_dev contains the device ID - // Together they uniquely identify a file - buffer_id.file_id = file_stat.st_ino; - std::filesystem::path p(file_name); - buffer_id.mtime = getLastModifiedNs(p); - } - return buffer_id; - } - - static BufferID VectorID(const std::string &file_name, uint32_t offset, - uint32_t length) { - BufferID buffer_id{}; - buffer_id.type = TYPE::VECTOR; - buffer_id.file_name = file_name; - struct stat file_stat; - if (stat(file_name.c_str(), &file_stat) == 0) { - buffer_id.file_id = file_stat.st_ino; - std::filesystem::path p(file_name); - buffer_id.mtime = getLastModifiedNs(p); - } - buffer_id.pos.vector.offset = offset; - buffer_id.pos.vector.length = length; - return buffer_id; - } - - explicit BufferID() = default; - - // Type of the file backing this buffer - TYPE type{UNKNOWN}; - - // Name of the file backing this buffer - std::string file_name{}; - - // Unique file id - uint64_t file_id{}; - - long mtime{}; - - // To identify which part of the backing file should be loaded into the buffer - Position pos{}; - - - // Get the forward ID - const inline struct ParquetPos &parquet() const { - return pos.forward; - } - - - // Get the vector ID - const inline struct VectorPos &vector() const { - return pos.vector; - } - - - // Get debug string - const std::string to_string() const { - std::string msg{"Buffer["}; - if (type == TYPE::PARQUET) { - msg += "parquet: " + file_name + "[" + std::to_string(file_id) + "]" + - ", column: " + std::to_string(parquet().column) + - ", row_group: " + std::to_string(parquet().row_group); - } else if (type == TYPE::VECTOR) { - msg += "vector: " + file_name + "[" + std::to_string(file_id) + "]" + - ", offset: " + std::to_string(vector().offset); - } else { - msg += "unknown"; - } - msg += ", mtime: " + std::to_string(mtime); - msg += "]"; - return msg; - } -}; - - -// Thread-safe LRU buffer implementation. -class BufferManager : public Singleton { - friend BufferHandle; - - public: - void init(uint64_t limit, uint32_t num_shards = 1); - - BufferHandle acquire(BufferID &buffer_id); - - std::unique_ptr acquire_ptr(BufferID &buffer_id); - - uint64_t total_size_in_bytes() const; - - void cleanup(); - - ~BufferManager(); - - private: - struct BufferContext; - - class BufferPool; - - // Custom deleter for Arrow buffer that automatically notifies us when the - // buffer is no longer referenced by Arrow - struct ArrowBufferDeleter { - explicit ArrowBufferDeleter(BufferContext *c); - BufferContext *context; - // Only reduces the reference count but does not actually release the - // buffer, since the buffer memory is managed by the BufferManager. - void operator()(arrow::Buffer *); - }; - - std::vector pools_; -}; - - -class BufferHandle { - public: - typedef std::unique_ptr Pointer; - - explicit BufferHandle(BufferManager::BufferContext *context = nullptr); - BufferHandle(const BufferHandle &) = delete; - BufferHandle(BufferHandle &&) = default; - BufferHandle &operator=(const BufferHandle &) = delete; - BufferHandle &operator=(BufferHandle &&) = default; - - - ~BufferHandle(); - - - // Pin parquet data in memory by allocating arrow buffers of appropriate size - // and reading data from the backing file. - // The lifecycle of the allocated memory is automatically managed through - // shared pointers. The buffers are guaranteed to be held until they are not - // referenced. - // Returns a pointer to the loaded ChunkedArray in Arrow format. - std::shared_ptr pin_parquet_data(); - - - // Pin vector data in memory by allocating a buffer of appropriate size and - // loading data from the backing file. - // The memory is guaranteed to be held until unpin() is called. The caller - // must call unpin() to release the memory when it is no longer needed. - // Returns a raw memory address. - void *pin_vector_data(); - - - // Reduce the reference count for this vector buffer. - // Returns true if this was the last reference. - // When reference count is zero, the buffer is moved to the eviction list and - // becomes eligible for removal under memory pressure. - bool unpin_vector_data(); - - - // Get the current reference count. - uint32_t references() const; - - - // Get the buffer size. - uint32_t size() const; - - - private: - using BufferContext = BufferManager::BufferContext; - using BufferPool = BufferManager::BufferPool; - - BufferContext *context_; - BufferPool *pool_; -}; - - -} // namespace ailego - - -} // namespace zvec diff --git a/src/include/zvec/ailego/buffer/parquet_hash_table.h b/src/include/zvec/ailego/buffer/parquet_hash_table.h index f4fb4f6..4db1a8f 100644 --- a/src/include/zvec/ailego/buffer/parquet_hash_table.h +++ b/src/include/zvec/ailego/buffer/parquet_hash_table.h @@ -114,7 +114,7 @@ class ParquetBufferPool { ParquetBufferPool *pool; ParquetBufferID id; // Only reduces the reference count but does not actually release the - // buffer, since the buffer memory is managed by the BufferManager. + // buffer, since the buffer memory is managed by the ParquetBufferPool. void operator()(arrow::Buffer *) { return; } diff --git a/tests/ailego/buffer/buffer_manager_test.cc b/tests/ailego/buffer/buffer_manager_test.cc deleted file mode 100644 index 90c9f4e..0000000 --- a/tests/ailego/buffer/buffer_manager_test.cc +++ /dev/null @@ -1,134 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed 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. - -#include -#include -#include -#include -#include -#include "tests/test_util.h" - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-result" -#endif - -using namespace zvec::ailego; - - -const std::string working_dir{"./buffer_manager_dir/"}; -const std::string file_path_forward{working_dir + "test.forward_index"}; -const std::string file_path_vector{working_dir + "test.vector_index"}; - - -class BufferManagerTest : public testing::Test { - /***** Global initialization and cleanup - Start *****/ - public: - static void SetUpTestCase() { - zvec::test_util::RemoveTestPath(working_dir); - - if (!File::MakePath(working_dir)) { - LOG_ERROR("Failed to create working directory."); - return; - } - - File file_vector_index; - size_t file_vector_size = 16 * 1024 * 1024; - if (!file_vector_index.create(file_path_vector, file_vector_size)) { - LOG_ERROR("Failed to create vector index file."); - return; - } - // Populate vector file with number series - for (uint32_t i = 0; i < file_vector_size / sizeof(uint32_t); ++i) { - file_vector_index.write((void *)&i, sizeof(i)); - } - file_vector_index.close(); - - BufferManager::Instance().init(4 * 1024 * 1024, 1); - } - - static void TearDownTestCase() { - BufferManager::Instance().cleanup(); - zvec::test_util::RemoveTestPath(working_dir); - } - /***** Global initialization and cleanup - End *****/ - ; -}; - - -TEST_F(BufferManagerTest, READ_VECTOR_FILE) { - uint32_t size_4KB = 4 * 1024; - - auto read_and_verify_numbers = [&](uint32_t offset) { - BufferID id = BufferID::VectorID(file_path_vector, offset, size_4KB); - auto handle = BufferManager::Instance().acquire(id); - uint32_t *vector_data = (uint32_t *)handle.pin_vector_data(); - uint32_t num_start = offset / sizeof(uint32_t); - for (uint32_t i = 0; i < size_4KB / sizeof(uint32_t); i++) { - ASSERT_EQ(*(vector_data + i), num_start + i); - } - handle.unpin_vector_data(); - }; - - std::vector threads; - - // Read the same part concurrently - for (int i = 0; i < 10; ++i) { - threads.emplace_back(read_and_verify_numbers, 3 * size_4KB); - } - for (auto &thread : threads) { - thread.join(); - } - - { // Verify the reference count - BufferID id = BufferID::VectorID(file_path_vector, 3 * size_4KB, size_4KB); - auto handle = BufferManager::Instance().acquire(id); - handle.pin_vector_data(); - ASSERT_EQ(handle.references(), 1); - handle.unpin_vector_data(); - ASSERT_EQ(handle.references(), 0); - } - - threads.clear(); - // Read different parts concurrently - for (int i = 0; i < 30; ++i) { - threads.emplace_back(read_and_verify_numbers, i * size_4KB); - } - for (auto &thread : threads) { - thread.join(); - } - ASSERT_EQ(BufferManager::Instance().total_size_in_bytes(), 30 * 4 * 1024); - - { // Read a large chunk so that the buffer is full - BufferID id = - BufferID::VectorID(file_path_vector, 4 * 1024 * 1024, 4 * 1024 * 1024); - auto handle = BufferManager::Instance().acquire(id); - handle.pin_vector_data(); - handle.unpin_vector_data(); - } - - { // Trigger eviction - BufferID id = - BufferID::VectorID(file_path_vector, 8 * 1024 * 1024, 4 * 1024 * 1024); - auto handle = BufferManager::Instance().acquire(id); - handle.pin_vector_data(); - ASSERT_EQ(BufferManager::Instance().total_size_in_bytes(), 4 * 1024 * 1024); - handle.unpin_vector_data(); - ASSERT_EQ(handle.references(), 0); - } -} - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic pop -#endif \ No newline at end of file diff --git a/tests/core/algorithm/flat/flat_streamer_buffer_test.cc b/tests/core/algorithm/flat/flat_streamer_buffer_test.cc index e3fce1f..e1f2177 100644 --- a/tests/core/algorithm/flat/flat_streamer_buffer_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_buffer_test.cc @@ -4,7 +4,6 @@ #include #include #include -#include #include #include #include "tests/test_util.h" diff --git a/tests/core/algorithm/flat/flat_streamer_buffer_time_test.cc b/tests/core/algorithm/flat/flat_streamer_buffer_time_test.cc index d84614d..ddf04d5 100644 --- a/tests/core/algorithm/flat/flat_streamer_buffer_time_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_buffer_time_test.cc @@ -4,7 +4,6 @@ #include #include #include -#include #include #include #include diff --git a/tests/core/algorithm/flat/flat_streamer_test.cc b/tests/core/algorithm/flat/flat_streamer_test.cc index 9d2c7c2..b6e57a7 100644 --- a/tests/core/algorithm/flat/flat_streamer_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_test.cc @@ -19,7 +19,6 @@ #include #include #include -#include #include #include #include diff --git a/tests/core/algorithm/flat_sparse/flat_sparse_streamer_buffer_test.cc b/tests/core/algorithm/flat_sparse/flat_sparse_streamer_buffer_test.cc index f85b182..7aed408 100644 --- a/tests/core/algorithm/flat_sparse/flat_sparse_streamer_buffer_test.cc +++ b/tests/core/algorithm/flat_sparse/flat_sparse_streamer_buffer_test.cc @@ -19,7 +19,6 @@ #include #include #include -#include #include #include #include "tests/test_util.h" diff --git a/tests/core/algorithm/flat_sparse/flat_sparse_streamer_test.cc b/tests/core/algorithm/flat_sparse/flat_sparse_streamer_test.cc index 7ac2493..88b5e26 100644 --- a/tests/core/algorithm/flat_sparse/flat_sparse_streamer_test.cc +++ b/tests/core/algorithm/flat_sparse/flat_sparse_streamer_test.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include "tests/test_util.h" diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc index cd21ff9..17c3c87 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc @@ -5,7 +5,6 @@ #include #include #include -#include #include #include #include "tests/test_util.h" diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 0f590ea..8d3adbd 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -23,7 +23,6 @@ #include "zvec/core/framework/index_provider.h" #endif #include -#include "zvec/ailego/buffer/buffer_manager.h" #include "zvec/core/interface/index.h" #include "zvec/core/interface/index_factory.h" #include "zvec/core/interface/index_param.h" @@ -295,7 +294,6 @@ TEST(IndexInterface, BufferGeneral) { .with_fetch_vector(true) .with_ef_search(20) .build()); - // zvec::ailego::BufferManager::Instance().cleanup(); } diff --git a/tests/core/utility/buffer_storage_test.cpp b/tests/core/utility/buffer_storage_test.cpp index 3e2d68a..c32e1b9 100644 --- a/tests/core/utility/buffer_storage_test.cpp +++ b/tests/core/utility/buffer_storage_test.cpp @@ -14,7 +14,6 @@ #include #include -#include #include #include diff --git a/tests/core/utility/buffer_storage_write_test.cc b/tests/core/utility/buffer_storage_write_test.cc index a97a32c..894c68f 100644 --- a/tests/core/utility/buffer_storage_write_test.cc +++ b/tests/core/utility/buffer_storage_write_test.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include diff --git a/tests/core/utility/mmap_file_storage_test.cpp b/tests/core/utility/mmap_file_storage_test.cpp index b0436d8..ed7a9d5 100644 --- a/tests/core/utility/mmap_file_storage_test.cpp +++ b/tests/core/utility/mmap_file_storage_test.cpp @@ -15,7 +15,6 @@ #include #include #include -#include #include #include diff --git a/tests/db/index/segment/segment_test.cc b/tests/db/index/segment/segment_test.cc index 4c0da0c..b3cd2f9 100644 --- a/tests/db/index/segment/segment_test.cc +++ b/tests/db/index/segment/segment_test.cc @@ -30,7 +30,6 @@ #include #include #include -#include #include "db/common/file_helper.h" #include "db/index/common/delete_store.h" #include "db/index/common/id_map.h"