wehub-resource-sync
556 lines
18 KiB
C++
556 lines
18 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 "zvec/core/framework/index_flow.h"
|
|
#include "meta_segment_common.h"
|
|
|
|
using namespace std;
|
|
|
|
namespace zvec {
|
|
namespace core {
|
|
|
|
#define SEARCH_DENSE_BATCH(_FUNC_NAME) \
|
|
int _FUNC_NAME(const void *query, const IndexQueryMeta &qmeta, \
|
|
uint32_t count, Context::Pointer &context) const { \
|
|
if (streamer_) { \
|
|
if (reformer_) { \
|
|
std::string ovec; \
|
|
IndexQueryMeta ometa; \
|
|
int ret = reformer_->convert(query, qmeta, count, &ovec, &ometa); \
|
|
if (ret != 0) { \
|
|
return ret; \
|
|
} \
|
|
return streamer_->_FUNC_NAME(ovec.data(), ometa, count, \
|
|
context->context()); \
|
|
} else { \
|
|
return streamer_->_FUNC_NAME(query, qmeta, count, context->context()); \
|
|
} \
|
|
} else { \
|
|
return flow_._FUNC_NAME(query, qmeta, count, context->flow_context()); \
|
|
} \
|
|
}
|
|
|
|
#define SEARCH_DENSE(_FUNC_NAME) \
|
|
int _FUNC_NAME(const void *query, const IndexQueryMeta &qmeta, \
|
|
Context::Pointer &context) const { \
|
|
if (streamer_) { \
|
|
if (reformer_) { \
|
|
std::string ovec; \
|
|
IndexQueryMeta ometa; \
|
|
int ret = reformer_->convert(query, qmeta, &ovec, &ometa); \
|
|
if (ret != 0) { \
|
|
return ret; \
|
|
} \
|
|
return streamer_->_FUNC_NAME(ovec.data(), ometa, context->context()); \
|
|
} else { \
|
|
return streamer_->_FUNC_NAME(query, qmeta, context->context()); \
|
|
} \
|
|
} else { \
|
|
return flow_._FUNC_NAME(query, qmeta, context->flow_context()); \
|
|
} \
|
|
}
|
|
|
|
#define SEARCH_SPRASE_BATCH(_FUNC_NAME) \
|
|
int _FUNC_NAME(const uint32_t *sparse_count, const uint32_t *sparse_indices, \
|
|
const void *sparse_query, const IndexQueryMeta &qmeta, \
|
|
uint32_t count, Context::Pointer &context) const { \
|
|
if (streamer_) { \
|
|
if (reformer_) { \
|
|
LOG_ERROR("reformer not supported in sparse search"); \
|
|
return IndexError_Runtime; \
|
|
} else { \
|
|
return streamer_->_FUNC_NAME(sparse_count, sparse_indices, \
|
|
sparse_query, qmeta, count, \
|
|
context->context()); \
|
|
} \
|
|
} else { \
|
|
return flow_._FUNC_NAME(sparse_count, sparse_indices, sparse_query, \
|
|
qmeta, count, context->flow_context()); \
|
|
} \
|
|
}
|
|
|
|
#define SEARCH_SPARSE(_FUNC_NAME) \
|
|
int _FUNC_NAME(const uint32_t sparse_count, const uint32_t *sparse_indices, \
|
|
const void *sparse_query, const IndexQueryMeta &qmeta, \
|
|
Context::Pointer &context) const { \
|
|
if (streamer_) { \
|
|
if (reformer_) { \
|
|
LOG_ERROR("reformer not supported in sparse search"); \
|
|
return IndexError_Runtime; \
|
|
} else { \
|
|
return streamer_->_FUNC_NAME(sparse_count, sparse_indices, \
|
|
sparse_query, qmeta, context->context()); \
|
|
} \
|
|
} else { \
|
|
return flow_._FUNC_NAME(sparse_count, sparse_indices, sparse_query, \
|
|
qmeta, context->flow_context()); \
|
|
} \
|
|
}
|
|
|
|
class Flow {
|
|
public:
|
|
class Context {
|
|
public:
|
|
typedef std::unique_ptr<Context> Pointer;
|
|
|
|
Context(IndexContext::Pointer &ctx, IndexFlow::Context::Pointer &flow_ctx)
|
|
: ctx_(std::move(ctx)), flow_ctx_(std::move(flow_ctx)) {}
|
|
|
|
void set_debug_mode(bool debug_mode) {
|
|
ctx_ ? ctx_->set_debug_mode(debug_mode)
|
|
: flow_ctx_->set_debug_mode(debug_mode);
|
|
}
|
|
|
|
std::string debug_string() {
|
|
return ctx_ ? ctx_->debug_string() : flow_ctx_->debug_string();
|
|
}
|
|
|
|
void set_topk(uint32_t topk) {
|
|
ctx_ ? ctx_->set_topk(topk) : flow_ctx_->set_topk(topk);
|
|
}
|
|
|
|
template <typename T>
|
|
void set_filter(T &&func) {
|
|
ctx_ ? ctx_->set_filter(func) : flow_ctx_->set_filter(func);
|
|
}
|
|
|
|
const IndexDocumentList &result(void) const {
|
|
return ctx_ ? ctx_->result() : flow_ctx_->result();
|
|
}
|
|
|
|
const IndexDocumentList &result(size_t index) const {
|
|
return ctx_ ? ctx_->result(index) : flow_ctx_->result(index);
|
|
}
|
|
|
|
public:
|
|
friend class Flow;
|
|
|
|
IndexFlow::Context::Pointer &flow_context(void) {
|
|
return flow_ctx_;
|
|
}
|
|
|
|
IndexContext::Pointer &context(void) {
|
|
return ctx_;
|
|
}
|
|
|
|
private:
|
|
IndexContext::Pointer ctx_;
|
|
IndexFlow::Context::Pointer flow_ctx_;
|
|
};
|
|
|
|
Context::Pointer create_context(void) const {
|
|
IndexContext::Pointer ctx;
|
|
IndexFlow::Context::Pointer flow_ctx;
|
|
if (streamer_) {
|
|
ctx = streamer_->create_context();
|
|
} else {
|
|
flow_ctx = flow_.create_context();
|
|
}
|
|
return Context::Pointer(new (std::nothrow) Context(ctx, flow_ctx));
|
|
}
|
|
|
|
int set_container(const std::string &name, const ailego::Params ¶ms) {
|
|
return flow_.set_storage(name, params);
|
|
}
|
|
|
|
int load_taglists(const std::string &path) {
|
|
// load tag lists
|
|
auto storage = IndexFactory::CreateStorage("MMapFileReadStorage");
|
|
|
|
int ret = storage->open(path, false);
|
|
if (ret != 0) {
|
|
LOG_ERROR("Failed to load index with storage %s",
|
|
storage->name().c_str());
|
|
return ret;
|
|
}
|
|
|
|
auto segment_taglist_header = storage->get(TAGLIST_HEADER_SEGMENT_NAME);
|
|
if (!segment_taglist_header) {
|
|
LOG_INFO("No Tag Lists Found!");
|
|
|
|
return 0;
|
|
}
|
|
|
|
TagListHeader taglist_header;
|
|
void *data_ptr;
|
|
if (segment_taglist_header->read(0, (const void **)(&data_ptr),
|
|
sizeof(TagListHeader)) !=
|
|
sizeof(TagListHeader)) {
|
|
LOG_ERROR("Read tag list meta failed");
|
|
return IndexError_ReadData;
|
|
}
|
|
|
|
memcpy(&taglist_header, data_ptr, sizeof(TagListHeader));
|
|
|
|
auto segment_taglist_key = storage->get(TAGLIST_KEY_SEGMENT_NAME);
|
|
if (!segment_taglist_key) {
|
|
LOG_ERROR("IndexStorage get segment %s failed",
|
|
TAGLIST_KEY_SEGMENT_NAME.c_str());
|
|
return IndexError_InvalidValue;
|
|
}
|
|
|
|
size_t offset = 0;
|
|
for (size_t i = 0; i < taglist_header.num_vecs; ++i) {
|
|
if (segment_taglist_key->read(offset, (const void **)(&data_ptr),
|
|
sizeof(uint64_t)) != sizeof(uint64_t)) {
|
|
LOG_ERROR("Read tag list key failed");
|
|
return IndexError_ReadData;
|
|
}
|
|
|
|
uint64_t key = *reinterpret_cast<const uint64_t *>(data_ptr);
|
|
tag_key_list_.push_back(key);
|
|
|
|
offset += sizeof(uint64_t);
|
|
}
|
|
|
|
auto segment_taglist_data = storage->get(TAGLIST_DATA_SEGMENT_NAME);
|
|
if (!segment_taglist_data) {
|
|
LOG_ERROR("IndexStorage get segment %s failed",
|
|
TAGLIST_DATA_SEGMENT_NAME.c_str());
|
|
return IndexError_InvalidValue;
|
|
}
|
|
|
|
std::vector<uint64_t> taglist_offsets;
|
|
offset = 0;
|
|
for (size_t i = 0; i < taglist_header.num_vecs; ++i) {
|
|
if (segment_taglist_data->read(offset, (const void **)(&data_ptr),
|
|
sizeof(uint64_t)) != sizeof(uint64_t)) {
|
|
LOG_ERROR("Read tag list data failed");
|
|
return IndexError_ReadData;
|
|
}
|
|
|
|
uint64_t tag_offset = *reinterpret_cast<const uint64_t *>(data_ptr);
|
|
taglist_offsets.push_back(tag_offset);
|
|
|
|
offset += sizeof(uint64_t);
|
|
}
|
|
|
|
offset = taglist_header.num_vecs * sizeof(uint64_t);
|
|
for (size_t i = 0; i < taglist_header.num_vecs; ++i) {
|
|
if (segment_taglist_data->read(offset, (const void **)(&data_ptr),
|
|
sizeof(uint64_t)) != sizeof(uint64_t)) {
|
|
LOG_ERROR("Read tag list data failed");
|
|
return IndexError_ReadData;
|
|
}
|
|
offset += sizeof(uint64_t);
|
|
|
|
uint64_t tag_count = *reinterpret_cast<const uint64_t *>(data_ptr);
|
|
|
|
if (segment_taglist_data->read(offset, (const void **)(&data_ptr),
|
|
tag_count * sizeof(uint64_t)) !=
|
|
tag_count * sizeof(uint64_t)) {
|
|
LOG_ERROR("Read tag list data failed");
|
|
return IndexError_ReadData;
|
|
}
|
|
offset += tag_count * sizeof(uint64_t);
|
|
|
|
std::vector<uint64_t> tag_list;
|
|
for (size_t j = 0; j < tag_count; ++j) {
|
|
uint64_t tag_id = *(reinterpret_cast<const uint64_t *>(data_ptr) + j);
|
|
tag_list.push_back(tag_id);
|
|
}
|
|
|
|
// order tags
|
|
sort(tag_list.begin(), tag_list.end());
|
|
|
|
id_to_tags_list_.push_back(std::move(tag_list));
|
|
}
|
|
|
|
storage->cleanup();
|
|
storage = nullptr;
|
|
|
|
return 0;
|
|
}
|
|
|
|
int load(const std::string &path) {
|
|
int ret = load_taglists(path);
|
|
if (ret != 0) {
|
|
LOG_ERROR("Failed to load tag lists");
|
|
return ret;
|
|
}
|
|
|
|
if (streamer_) {
|
|
stg_ = IndexFactory::CreateStorage("MMapFileStorage");
|
|
if (!stg_) {
|
|
return IndexError_NoExist;
|
|
}
|
|
ailego::Params params;
|
|
params.set("proxima.mmap_file.storage.memory_warmup", true);
|
|
ret = stg_->init(params);
|
|
if (ret != 0) {
|
|
return ret;
|
|
}
|
|
ret = stg_->open(path, true);
|
|
if (ret != 0) {
|
|
return ret;
|
|
}
|
|
|
|
if (!inited_) {
|
|
IndexMeta meta;
|
|
ret = IndexHelper::DeserializeFromStorage(stg_.get(), &meta);
|
|
if (ret != 0) {
|
|
LOG_ERROR("Failed to get IndexMeta from Storage");
|
|
return ret;
|
|
}
|
|
ret = streamer_->init(meta, searcher_params_);
|
|
if (ret != 0) {
|
|
return ret;
|
|
}
|
|
|
|
if (!meta.reformer_name().empty()) {
|
|
reformer_ = IndexFactory::CreateReformer(meta.reformer_name());
|
|
if (!reformer_) {
|
|
LOG_ERROR("Failed to create reformer %s",
|
|
meta.reformer_name().c_str());
|
|
return IndexError_NoExist;
|
|
}
|
|
reformer_->init(meta.reformer_params());
|
|
}
|
|
}
|
|
|
|
return streamer_->open(stg_);
|
|
} else {
|
|
return flow_.load(path);
|
|
}
|
|
}
|
|
|
|
int unload(void) {
|
|
if (streamer_) {
|
|
streamer_->close();
|
|
return stg_->close();
|
|
} else {
|
|
return flow_.unload();
|
|
}
|
|
}
|
|
|
|
int set_searcher(const std::string &name, const ailego::Params ¶ms) {
|
|
//! If the searcher is streamer, create it
|
|
streamer_ = IndexFactory::CreateStreamer(name);
|
|
if (!streamer_) {
|
|
return flow_.set_searcher(name, params);
|
|
}
|
|
searcher_params_ = params;
|
|
return 0;
|
|
}
|
|
|
|
int set_searcher(IndexStreamer::Pointer streamer) {
|
|
streamer_ = streamer;
|
|
|
|
inited_ = true;
|
|
|
|
return 0;
|
|
}
|
|
|
|
const std::vector<std::vector<uint64_t>> &id_to_tags_list() const {
|
|
return id_to_tags_list_;
|
|
}
|
|
|
|
const std::vector<uint64_t> &tag_key_list() const {
|
|
return tag_key_list_;
|
|
}
|
|
|
|
SEARCH_DENSE_BATCH(search_impl);
|
|
SEARCH_DENSE(search_impl);
|
|
SEARCH_DENSE_BATCH(search_bf_impl);
|
|
SEARCH_DENSE(search_bf_impl);
|
|
|
|
private:
|
|
IndexFlow flow_{};
|
|
|
|
IndexStreamer::Pointer streamer_{};
|
|
IndexReformer::Pointer reformer_{};
|
|
|
|
bool inited_{false};
|
|
|
|
IndexStorage::Pointer stg_{};
|
|
ailego::Params searcher_params_{};
|
|
std::vector<std::vector<uint64_t>> id_to_tags_list_;
|
|
std::vector<uint64_t> tag_key_list_;
|
|
};
|
|
|
|
class SparseFlow {
|
|
public:
|
|
class Context {
|
|
public:
|
|
typedef std::unique_ptr<Context> Pointer;
|
|
|
|
Context(IndexContext::Pointer &ctx,
|
|
IndexSparseFlow::Context::Pointer &flow_ctx)
|
|
: ctx_(std::move(ctx)), flow_ctx_(std::move(flow_ctx)) {}
|
|
|
|
void set_debug_mode(bool debug_mode) {
|
|
ctx_ ? ctx_->set_debug_mode(debug_mode)
|
|
: flow_ctx_->set_debug_mode(debug_mode);
|
|
}
|
|
|
|
std::string debug_string() {
|
|
return ctx_ ? ctx_->debug_string() : flow_ctx_->debug_string();
|
|
}
|
|
|
|
template <typename T>
|
|
void set_filter(T &&func) {
|
|
ctx_ ? ctx_->set_filter(func) : flow_ctx_->set_filter(func);
|
|
}
|
|
|
|
void set_topk(uint32_t topk) {
|
|
ctx_ ? ctx_->set_topk(topk) : flow_ctx_->set_topk(topk);
|
|
}
|
|
|
|
const IndexDocumentList &result(void) const {
|
|
return ctx_ ? ctx_->result() : flow_ctx_->result();
|
|
}
|
|
|
|
const IndexDocumentList &result(size_t index) const {
|
|
return ctx_ ? ctx_->result(index) : flow_ctx_->result(index);
|
|
}
|
|
|
|
private:
|
|
friend class SparseFlow;
|
|
|
|
IndexSparseFlow::Context::Pointer &flow_context(void) {
|
|
return flow_ctx_;
|
|
}
|
|
|
|
IndexContext::Pointer &context(void) {
|
|
return ctx_;
|
|
}
|
|
|
|
|
|
private:
|
|
IndexContext::Pointer ctx_;
|
|
IndexSparseFlow::Context::Pointer flow_ctx_;
|
|
};
|
|
|
|
Context::Pointer create_context(void) const {
|
|
IndexContext::Pointer ctx;
|
|
IndexSparseFlow::Context::Pointer flow_ctx;
|
|
if (streamer_) {
|
|
ctx = streamer_->create_context();
|
|
} else {
|
|
flow_ctx = flow_.create_context();
|
|
}
|
|
return Context::Pointer(new (std::nothrow) Context(ctx, flow_ctx));
|
|
}
|
|
|
|
int set_container(const std::string &name, const ailego::Params ¶ms) {
|
|
return flow_.set_storage(name, params);
|
|
}
|
|
|
|
int load(const std::string &path) {
|
|
if (streamer_) {
|
|
stg_ = IndexFactory::CreateStorage("MMapFileStorage");
|
|
if (!stg_) {
|
|
return IndexError_NoExist;
|
|
}
|
|
ailego::Params params;
|
|
params.set("proxima.mmap_file.storage.memory_warmup", true);
|
|
int ret = stg_->init(params);
|
|
if (ret != 0) {
|
|
return ret;
|
|
}
|
|
ret = stg_->open(path, true);
|
|
if (ret != 0) {
|
|
return ret;
|
|
}
|
|
|
|
if (!inited_) {
|
|
IndexMeta meta;
|
|
ret = IndexHelper::DeserializeFromStorage(stg_.get(), &meta);
|
|
if (ret != 0) {
|
|
LOG_ERROR("Failed to get IndexMeta from Storage");
|
|
return ret;
|
|
}
|
|
|
|
ret = streamer_->init(meta, searcher_params_);
|
|
if (ret != 0) {
|
|
return ret;
|
|
}
|
|
|
|
if (!meta.reformer_name().empty()) {
|
|
reformer_ = IndexFactory::CreateReformer(meta.reformer_name());
|
|
if (!reformer_) {
|
|
LOG_ERROR("Failed to create reformer %s",
|
|
meta.reformer_name().c_str());
|
|
return IndexError_NoExist;
|
|
}
|
|
reformer_->init(meta.reformer_params());
|
|
}
|
|
}
|
|
|
|
return streamer_->open(stg_);
|
|
} else {
|
|
return flow_.load(path);
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int unload(void) {
|
|
if (streamer_) {
|
|
streamer_->close();
|
|
return stg_->close();
|
|
} else {
|
|
return flow_.unload();
|
|
}
|
|
}
|
|
|
|
int set_searcher(const std::string &name, const ailego::Params ¶ms) {
|
|
//! If the searcher is streamer, create it
|
|
streamer_ = IndexFactory::CreateStreamer(name);
|
|
if (!streamer_) {
|
|
return flow_.set_searcher(name, params);
|
|
}
|
|
searcher_params_ = params;
|
|
return 0;
|
|
}
|
|
|
|
int set_searcher(IndexStreamer::Pointer streamer) {
|
|
streamer_ = streamer;
|
|
|
|
inited_ = true;
|
|
|
|
return 0;
|
|
}
|
|
|
|
const std::vector<std::vector<uint64_t>> &id_to_tags_list() const {
|
|
return id_to_tags_list_;
|
|
}
|
|
|
|
const std::vector<uint64_t> &tag_key_list() const {
|
|
return tag_key_list_;
|
|
}
|
|
|
|
SEARCH_SPRASE_BATCH(search_impl);
|
|
SEARCH_SPARSE(search_impl);
|
|
SEARCH_SPRASE_BATCH(search_bf_impl);
|
|
SEARCH_SPARSE(search_bf_impl);
|
|
|
|
private:
|
|
IndexSparseFlow flow_{};
|
|
|
|
IndexStreamer::Pointer streamer_{};
|
|
IndexReformer::Pointer reformer_{};
|
|
|
|
bool inited_{false};
|
|
|
|
IndexStorage::Pointer stg_{};
|
|
ailego::Params searcher_params_{};
|
|
std::vector<std::vector<uint64_t>> id_to_tags_list_;
|
|
std::vector<uint64_t> tag_key_list_;
|
|
};
|
|
|
|
} // namespace core
|
|
} // namespace zvec
|