chore: rm buffer manager (#437)
This commit is contained in:
ZeFeng Yin
GitHub
parent
8e8bb81db0
commit
8ce8e3e228
|
|
@ -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 <atomic>
|
||||
#include <mutex>
|
||||
#include <ailego/pattern/defer.h>
|
||||
#include <arrow/io/api.h>
|
||||
#include <parquet/arrow/reader.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/internal/platform.h>
|
||||
#include <zvec/ailego/logger/logger.h>
|
||||
|
||||
#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 <arrow/api.h>
|
||||
|
||||
#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<int>{}(static_cast<int>(buffer_id.type));
|
||||
hash = hash ^ (std::hash<uint64_t>{}(buffer_id.file_id));
|
||||
if (buffer_id.type == BufferID::TYPE::PARQUET) {
|
||||
hash = hash * 31 + std::hash<int>{}(buffer_id.parquet().column);
|
||||
hash = hash * 31 + std::hash<int>{}(buffer_id.parquet().row_group);
|
||||
} else if (buffer_id.type == BufferID::TYPE::VECTOR) {
|
||||
hash = hash * 31 + std::hash<uint32_t>{}(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<BufferManager::BufferContext> 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<uint32_t> refs_buf{0};
|
||||
|
||||
// The number of external references to this context (via BufferHandle)
|
||||
std::atomic<uint32_t> refs_context{0};
|
||||
|
||||
BufferPool *pool{nullptr};
|
||||
|
||||
// A shared pointer to the buffers allocated for arrow parquet data
|
||||
std::shared_ptr<arrow::ChunkedArray> arrow{nullptr};
|
||||
|
||||
// Guard original arrow buffers to prevent premature deletion
|
||||
std::vector<std::shared_ptr<arrow::Buffer>> 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<arrow::io::RandomAccessFile> input;
|
||||
const auto &file_name = id.file_name;
|
||||
ARROW_ASSIGN_OR_RAISE(input, arrow::io::ReadableFile::Open(file_name));
|
||||
|
||||
// Open reader
|
||||
std::unique_ptr<parquet::arrow::FileReader> 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<arrow::Buffer> 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<std::mutex> lock(mutex_context_);
|
||||
if (auto iter = table_.find(id); iter != table_.end()) {
|
||||
return iter->second.get();
|
||||
}
|
||||
auto [iter, _] =
|
||||
table_.emplace(id, std::make_unique<BufferContext>(id, this));
|
||||
return iter->second.get();
|
||||
}
|
||||
|
||||
|
||||
void try_release_context_locked(BufferContext *context) {
|
||||
if (context->refs_context.load() != 0) {
|
||||
return;
|
||||
}
|
||||
std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> 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<std::mutex> lock(mutex_context_);
|
||||
LRU_insert(context);
|
||||
}
|
||||
|
||||
|
||||
void LRU_remove_locked(BufferContext *context) {
|
||||
std::lock_guard<std::mutex> 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<BufferID, BufferContext::Pointer, IDHash, IDEqual>;
|
||||
|
||||
uint64_t limit_;
|
||||
std::atomic<uint64_t> usage_{0};
|
||||
std::atomic<uint64_t> 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<arrow::ChunkedArray> 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<BufferHandle> 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<BufferHandle>(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
|
||||
|
|
@ -14,7 +14,6 @@
|
|||
#include "db/common/global_resource.h"
|
||||
#include <mutex>
|
||||
#include <zvec/ailego/buffer/block_eviction_queue.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/db/config.h>
|
||||
|
||||
namespace zvec {
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@
|
|||
#include <arrow/result.h>
|
||||
#include <arrow/status.h>
|
||||
#include <parquet/arrow/reader.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/buffer/parquet_hash_table.h>
|
||||
#include <zvec/ailego/logger/logger.h>
|
||||
#include "db/index/storage/store_helper.h"
|
||||
|
|
|
|||
|
|
@ -23,7 +23,6 @@
|
|||
#include <arrow/ipc/reader.h>
|
||||
#include <arrow/util/async_generator.h>
|
||||
#include <parquet/arrow/reader.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/db/status.h>
|
||||
#include "base_forward_store.h"
|
||||
|
||||
|
|
@ -123,9 +122,6 @@ class BufferPoolForwardStore
|
|||
|
||||
/// Number of rows in each row group
|
||||
std::vector<int64_t> row_group_row_nums_;
|
||||
|
||||
/// Buffer manager for caching data
|
||||
std::shared_ptr<ailego::BufferManager> buffer_manager_;
|
||||
};
|
||||
|
||||
} // namespace zvec
|
||||
|
|
@ -16,7 +16,6 @@
|
|||
|
||||
#include <arrow/ipc/reader.h>
|
||||
#include <parquet/arrow/reader.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/buffer/parquet_hash_table.h>
|
||||
#include "db/common/constants.h"
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <sys/stat.h>
|
||||
#include <chrono>
|
||||
#include <cstdint>
|
||||
#include <filesystem>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
#include <zvec/ailego/io/file.h>
|
||||
#include <zvec/ailego/pattern/singleton.h>
|
||||
|
||||
namespace arrow {
|
||||
class ChunkedArray;
|
||||
class Array;
|
||||
class DataType;
|
||||
class Scalar;
|
||||
template <typename T>
|
||||
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<std::uint64_t>(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<BufferManager> {
|
||||
friend BufferHandle;
|
||||
|
||||
public:
|
||||
void init(uint64_t limit, uint32_t num_shards = 1);
|
||||
|
||||
BufferHandle acquire(BufferID &buffer_id);
|
||||
|
||||
std::unique_ptr<BufferHandle> 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<BufferPool *> pools_;
|
||||
};
|
||||
|
||||
|
||||
class BufferHandle {
|
||||
public:
|
||||
typedef std::unique_ptr<BufferHandle> 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<arrow::ChunkedArray> 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
|
||||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 <cstdint>
|
||||
#include <thread>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/logger/logger.h>
|
||||
#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<std::thread> 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
|
||||
|
|
@ -4,7 +4,6 @@
|
|||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@
|
|||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/utility/file_helper.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@
|
|||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/encoding/json/mod_json.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@
|
|||
#include <ailego/utility/memory_helper.h>
|
||||
#include <algorithm/flat_sparse/flat_sparse_utility.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@
|
|||
#include <ailego/utility/memory_helper.h>
|
||||
#include <algorithm/flat_sparse/flat_sparse_utility.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@
|
|||
#include <ailego/utility/memory_helper.h>
|
||||
#include <algorithm/hnsw/hnsw_params.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
|
|
|||
|
|
@ -23,7 +23,6 @@
|
|||
#include "zvec/core/framework/index_provider.h"
|
||||
#endif
|
||||
#include <zvec/ailego/buffer/block_eviction_queue.h>
|
||||
#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();
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,6 @@
|
|||
|
||||
#include <iostream>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include <zvec/core/framework/index_helper.h>
|
||||
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@
|
|||
#include <vector>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/block_eviction_queue.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/ailego/io/file.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include <zvec/core/framework/index_helper.h>
|
||||
|
|
|
|||
|
|
@ -15,7 +15,6 @@
|
|||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include <zvec/core/framework/index_helper.h>
|
||||
|
||||
|
|
|
|||
|
|
@ -30,7 +30,6 @@
|
|||
#include <arrow/table.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/buffer/block_eviction_queue.h>
|
||||
#include <zvec/ailego/buffer/buffer_manager.h>
|
||||
#include "db/common/file_helper.h"
|
||||
#include "db/index/common/delete_store.h"
|
||||
#include "db/index/common/id_map.h"
|
||||
|
|
|
|||
Loading…
Reference in New Issue