wehub-resource-sync
378 lines
10 KiB
C++
378 lines
10 KiB
C++
// 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 <iostream>
|
|
#include <zvec/ailego/io/mmap_file.h>
|
|
#include "zvec/core/framework/index_meta.h"
|
|
#include "vecs_common.h"
|
|
|
|
namespace zvec {
|
|
namespace core {
|
|
|
|
class VecsReader {
|
|
public:
|
|
VecsReader()
|
|
: mmap_file_(),
|
|
index_meta_(),
|
|
num_vecs_(0),
|
|
vector_base_(nullptr),
|
|
key_base_(nullptr),
|
|
sparse_base_meta_{nullptr},
|
|
sparse_base_data_{nullptr},
|
|
partition_base_{nullptr},
|
|
taglist_base_meta_{nullptr},
|
|
taglist_base_data_{nullptr},
|
|
taglist_size_{0} {}
|
|
|
|
void set_metric(const std::string &name, const ailego::Params ¶ms) {
|
|
index_meta_.set_metric(name, 0, params);
|
|
}
|
|
|
|
bool load(const std::string &fname) {
|
|
return load(fname.c_str());
|
|
}
|
|
|
|
bool load(const char *fname) {
|
|
if (!fname) {
|
|
std::cerr << "Load fname is nullptr" << std::endl;
|
|
return false;
|
|
}
|
|
if (!mmap_file_.open(fname, true)) {
|
|
std::cerr << "Open file error: " << fname << std::endl;
|
|
return false;
|
|
}
|
|
|
|
return load();
|
|
}
|
|
|
|
bool load() {
|
|
const VecsHeader *header =
|
|
reinterpret_cast<const VecsHeader *>(mmap_file_.region());
|
|
// check
|
|
num_vecs_ = header->num_vecs;
|
|
|
|
// deserialize
|
|
bool bret = index_meta_.deserialize(header->meta_buf(), header->meta_size);
|
|
if (!bret) {
|
|
std::cerr << "deserialize index meta error." << std::endl;
|
|
return false;
|
|
}
|
|
|
|
const char *data_base_ptr =
|
|
reinterpret_cast<const char *>(header->meta_buf()) + header->meta_size;
|
|
|
|
vector_base_ = reinterpret_cast<const char *>(data_base_ptr);
|
|
key_base_ = reinterpret_cast<const uint64_t *>(
|
|
vector_base_ + num_vecs_ * index_meta_.element_size());
|
|
|
|
if (header->sparse_offset != -1LLU) {
|
|
sparse_base_meta_ = data_base_ptr + header->sparse_offset;
|
|
sparse_base_data_ = sparse_base_meta_ + num_vecs_ * sizeof(uint64_t);
|
|
}
|
|
|
|
if (header->partition_offset != -1LLU) {
|
|
partition_base_ = reinterpret_cast<const uint32_t *>(
|
|
data_base_ptr + header->partition_offset);
|
|
}
|
|
|
|
if (header->taglist_offset != -1LLU) {
|
|
taglist_base_meta_ = data_base_ptr + header->taglist_offset;
|
|
taglist_base_data_ = taglist_base_meta_ + num_vecs_;
|
|
taglist_size_ = header->taglist_size;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
size_t num_vecs() const {
|
|
return num_vecs_;
|
|
}
|
|
|
|
const void *vector_base() const {
|
|
return vector_base_;
|
|
}
|
|
|
|
const uint64_t *key_base() const {
|
|
return key_base_;
|
|
}
|
|
|
|
const IndexMeta &index_meta() const {
|
|
return index_meta_;
|
|
}
|
|
|
|
uint64_t get_key(size_t index) const {
|
|
return key_base_[index];
|
|
}
|
|
|
|
const void *get_vector(size_t index) const {
|
|
return vector_base_ + index * index_meta_.element_size();
|
|
}
|
|
|
|
uint32_t get_sparse_count(size_t index) const {
|
|
auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t);
|
|
uint64_t sparse_offset = *((uint64_t *)sparse_data_meta);
|
|
uint32_t sparse_count = *((uint32_t *)(sparse_base_data_ + sparse_offset));
|
|
|
|
return sparse_count;
|
|
|
|
return 0;
|
|
}
|
|
|
|
const uint32_t *get_sparse_indices(size_t index) const {
|
|
auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t);
|
|
uint64_t sparse_offset = *((uint64_t *)sparse_data_meta);
|
|
uint32_t *sparse_indices =
|
|
(uint32_t *)(sparse_base_data_ + sparse_offset + sizeof(uint32_t));
|
|
|
|
return sparse_indices;
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
const void *get_sparse_data(size_t index) const {
|
|
auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t);
|
|
uint64_t sparse_offset = *((uint64_t *)sparse_data_meta);
|
|
uint32_t sparse_count = *((uint32_t *)(sparse_base_data_ + sparse_offset));
|
|
void *sparse_data =
|
|
(uint32_t *)(sparse_base_data_ + sparse_offset + sizeof(uint32_t) +
|
|
sparse_count * sizeof(uint32_t));
|
|
|
|
return sparse_data;
|
|
}
|
|
|
|
size_t get_total_sparse_count(void) const {
|
|
size_t total_sparse_count = 0;
|
|
for (size_t i = 0; i < num_vecs_; ++i) {
|
|
total_sparse_count += get_sparse_count(i);
|
|
}
|
|
|
|
return total_sparse_count;
|
|
}
|
|
|
|
bool has_taglist(void) const {
|
|
return taglist_base_meta_ != nullptr;
|
|
}
|
|
|
|
uint64_t get_taglist_count(size_t index) const {
|
|
if (!taglist_base_data_ || !taglist_base_meta_) {
|
|
return 0;
|
|
}
|
|
|
|
uint64_t taglist_count = *reinterpret_cast<const uint64_t *>(
|
|
taglist_base_data_ + taglist_base_meta_[index]);
|
|
return taglist_count;
|
|
}
|
|
|
|
const uint64_t *get_taglist(size_t index) const {
|
|
if (!taglist_base_data_ || !taglist_base_meta_) {
|
|
return nullptr;
|
|
}
|
|
|
|
return reinterpret_cast<const uint64_t *>(taglist_base_data_ +
|
|
taglist_base_meta_[index]) +
|
|
1;
|
|
}
|
|
|
|
const void *get_taglist_data(size_t &size) const {
|
|
size = taglist_size_;
|
|
|
|
return taglist_base_meta_;
|
|
}
|
|
|
|
private:
|
|
ailego::MMapFile mmap_file_;
|
|
IndexMeta index_meta_;
|
|
size_t num_vecs_;
|
|
const char *vector_base_;
|
|
const uint64_t *key_base_;
|
|
const char *sparse_base_meta_;
|
|
const char *sparse_base_data_;
|
|
const uint32_t *partition_base_;
|
|
const char *taglist_base_meta_;
|
|
const char *taglist_base_data_;
|
|
uint64_t taglist_size_;
|
|
};
|
|
|
|
class SparseVecsReader {
|
|
public:
|
|
SparseVecsReader()
|
|
: mmap_file_(),
|
|
index_meta_(),
|
|
num_vecs_(0),
|
|
key_base_(nullptr),
|
|
sparse_base_meta_(nullptr),
|
|
sparse_base_data_{nullptr},
|
|
partition_base_{nullptr},
|
|
taglist_base_meta_{nullptr},
|
|
taglist_base_data_{nullptr},
|
|
taglist_size_{0} {}
|
|
|
|
void set_metric(const std::string &name, const ailego::Params ¶ms) {
|
|
index_meta_.set_metric(name, 0, params);
|
|
}
|
|
|
|
bool load(const std::string &fname) {
|
|
return load(fname.c_str());
|
|
}
|
|
|
|
|
|
bool load(const char *fname) {
|
|
if (!fname) {
|
|
std::cerr << "Load fname is nullptr" << std::endl;
|
|
return false;
|
|
}
|
|
if (!mmap_file_.open(fname, true)) {
|
|
std::cerr << "Open file error: " << fname << std::endl;
|
|
return false;
|
|
}
|
|
|
|
return load();
|
|
}
|
|
|
|
bool load() {
|
|
const VecsHeader *header =
|
|
reinterpret_cast<const VecsHeader *>(mmap_file_.region());
|
|
|
|
// check
|
|
num_vecs_ = header->num_vecs;
|
|
|
|
// deserialize
|
|
bool bret = index_meta_.deserialize(header->meta_buf(), header->meta_size);
|
|
if (!bret) {
|
|
std::cerr << "deserialize index meta error." << std::endl;
|
|
return false;
|
|
}
|
|
|
|
const char *data_base_ptr =
|
|
reinterpret_cast<const char *>(header->meta_buf()) + header->meta_size;
|
|
|
|
key_base_ = reinterpret_cast<const uint64_t *>(
|
|
reinterpret_cast<const char *>(header->meta_buf()) + header->meta_size);
|
|
sparse_base_meta_ = reinterpret_cast<const char *>(key_base_ + num_vecs_);
|
|
sparse_base_data_ = reinterpret_cast<const char *>(
|
|
sparse_base_meta_ + num_vecs_ * sizeof(uint64_t));
|
|
|
|
if (header->partition_offset != -1LLU) {
|
|
partition_base_ = reinterpret_cast<const uint32_t *>(
|
|
data_base_ptr + header->partition_offset);
|
|
}
|
|
|
|
if (header->taglist_offset != -1LLU) {
|
|
taglist_base_meta_ = data_base_ptr + header->taglist_offset;
|
|
taglist_base_data_ = taglist_base_meta_ + num_vecs_;
|
|
taglist_size_ = header->taglist_size;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
size_t num_vecs() const {
|
|
return num_vecs_;
|
|
}
|
|
|
|
const void *sparse_meta_base() const {
|
|
return sparse_base_meta_;
|
|
}
|
|
|
|
const uint64_t *key_base() const {
|
|
return key_base_;
|
|
}
|
|
|
|
const IndexMeta &index_meta() const {
|
|
return index_meta_;
|
|
}
|
|
|
|
uint64_t get_key(size_t index) const {
|
|
return key_base_[index];
|
|
}
|
|
|
|
uint32_t get_sparse_count(size_t index) const {
|
|
auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t);
|
|
uint64_t sparse_offset = *((uint64_t *)sparse_data_meta);
|
|
uint32_t sparse_count = *((uint32_t *)(sparse_base_data_ + sparse_offset));
|
|
|
|
return sparse_count;
|
|
|
|
return 0;
|
|
}
|
|
|
|
const uint32_t *get_sparse_indices(size_t index) const {
|
|
auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t);
|
|
uint64_t sparse_offset = *((uint64_t *)sparse_data_meta);
|
|
uint32_t *sparse_indices =
|
|
(uint32_t *)(sparse_base_data_ + sparse_offset + sizeof(uint32_t));
|
|
|
|
return sparse_indices;
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
const void *get_sparse_data(size_t index) const {
|
|
auto sparse_data_meta = sparse_base_meta_ + index * sizeof(uint64_t);
|
|
uint64_t sparse_offset = *((uint64_t *)sparse_data_meta);
|
|
uint32_t sparse_count = *((uint32_t *)(sparse_base_data_ + sparse_offset));
|
|
void *sparse_data =
|
|
(uint32_t *)(sparse_base_data_ + sparse_offset + sizeof(uint32_t) +
|
|
sparse_count * sizeof(uint32_t));
|
|
|
|
return sparse_data;
|
|
}
|
|
|
|
size_t get_total_sparse_count(void) const {
|
|
size_t total_sparse_count = 0;
|
|
for (size_t i = 0; i < num_vecs_; ++i) {
|
|
total_sparse_count += get_sparse_count(i);
|
|
}
|
|
|
|
return total_sparse_count;
|
|
}
|
|
|
|
bool has_taglist(void) const {
|
|
return taglist_base_meta_ != nullptr;
|
|
}
|
|
|
|
uint64_t get_taglist_count(size_t index) const {
|
|
uint64_t taglist_count = *reinterpret_cast<const uint64_t *>(
|
|
taglist_base_data_ + taglist_base_meta_[index]);
|
|
return taglist_count;
|
|
}
|
|
|
|
const uint64_t *get_taglist(size_t index) const {
|
|
return reinterpret_cast<const uint64_t *>(taglist_base_data_ +
|
|
taglist_base_meta_[index]) +
|
|
1;
|
|
}
|
|
|
|
const void *get_taglist_data(size_t &size) const {
|
|
size = taglist_size_;
|
|
return taglist_base_meta_;
|
|
}
|
|
|
|
private:
|
|
ailego::MMapFile mmap_file_;
|
|
IndexMeta index_meta_;
|
|
size_t num_vecs_;
|
|
const uint64_t *key_base_;
|
|
const char *sparse_base_meta_;
|
|
const char *sparse_base_data_;
|
|
const uint32_t *partition_base_;
|
|
const char *taglist_base_meta_;
|
|
const char *taglist_base_data_;
|
|
uint64_t taglist_size_;
|
|
};
|
|
|
|
} // namespace core
|
|
} // namespace zvec
|