chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,59 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
|
||||
|
||||
cc_binary(
|
||||
NAME txt2vecs
|
||||
STRICT PACKED
|
||||
SRCS txt2vecs.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags core_framework zvec_ailego
|
||||
)
|
||||
|
||||
cc_binary(
|
||||
NAME local_builder
|
||||
STRICT PACKED
|
||||
SRCS local_builder.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface core_knn_diskann
|
||||
)
|
||||
|
||||
cc_binary(
|
||||
NAME recall
|
||||
STRICT PACKED
|
||||
SRCS recall.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann
|
||||
)
|
||||
|
||||
cc_binary(
|
||||
NAME bench
|
||||
STRICT PACKED
|
||||
SRCS bench.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann
|
||||
)
|
||||
|
||||
|
||||
cc_binary(
|
||||
NAME recall_original
|
||||
STRICT PACKED
|
||||
SRCS recall_original.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann
|
||||
)
|
||||
|
||||
cc_binary(
|
||||
NAME bench_original
|
||||
STRICT PACKED
|
||||
SRCS bench_original.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann
|
||||
)
|
||||
|
||||
cc_binary(
|
||||
NAME local_builder_original
|
||||
STRICT PACKED
|
||||
SRCS local_builder_original.cc
|
||||
INCS ${PROJECT_ROOT_DIR}/src/core/
|
||||
LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface core_knn_diskann
|
||||
)
|
||||
@@ -0,0 +1,150 @@
|
||||
|
||||
# Benchmarking scripts
|
||||
|
||||
This directory contains benchmarking scripts and reproducing steps.
|
||||
|
||||
## COHERE experiments
|
||||
|
||||
### Getting COHERE Data
|
||||
|
||||
Please download the COHERE 10M dataset to cohere_large_10m as follows:
|
||||
|
||||
```bash
|
||||
... ...
|
||||
neighbors.parquet
|
||||
shuffle_train-00-of-10.parquet
|
||||
shuffle_train-01-of-10.parquet
|
||||
shuffle_train-02-of-10.parquet
|
||||
shuffle_train-03-of-10.parquet
|
||||
shuffle_train-04-of-10.parquet
|
||||
shuffle_train-05-of-10.parquet
|
||||
shuffle_train-06-of-10.parquet
|
||||
shuffle_train-07-of-10.parquet
|
||||
shuffle_train-08-of-10.parquet
|
||||
shuffle_train-09-of-10.parquet
|
||||
scalar_labels.parquet
|
||||
test.parquet
|
||||
```
|
||||
|
||||
For convenience, we prepared a docker image with cohere bench datasets: registry.cn-hongkong.cr.aliyuncs.com/zvec/cohere-bench-data.
|
||||
|
||||
You can run a container as follows:
|
||||
|
||||
```bash
|
||||
docker run -it --net=host -d -e DEBUG_MODE=true --user root --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /home/zvec/:/home/zvec/ -w /home/zvec --name=cohere_bench zvec-registry.cn-hongkong.cr.aliyuncs.com/zvec/cohere-bench-data:0.0.1 bash
|
||||
|
||||
docker exec -it cohere_bench bash
|
||||
```
|
||||
|
||||
The datasets locate at /tmp/cohere/
|
||||
|
||||
### Preparing Environment
|
||||
Clone code and init:
|
||||
```bash
|
||||
$ git clone git@github.com:alibaba/zvec.git
|
||||
$ cd zvec
|
||||
$ git submodule update --init
|
||||
```
|
||||
|
||||
Build source code:
|
||||
```
|
||||
$ cd /home/zvec/workspace/zvec
|
||||
$ mkdir build
|
||||
$ cd build
|
||||
$ cmake -DENABLE_SKYLAKE=ON -DCMAKE_BUILD_TYPE=Release ..
|
||||
```
|
||||
|
||||
### Converting Dataset
|
||||
Export vector data using python script:
|
||||
```bash
|
||||
$ mkdir 10m.output
|
||||
$ python3 convert_cohere_parquet.py
|
||||
```
|
||||
|
||||
Convert vector data to binary formatted file.
|
||||
```bash
|
||||
/home/zvec/workspace/zvec/bin/txt2vecs -input=cohere_train_vector_10m.txt --output=cohere_train_vector_10m.zvec.vecs --dimension=768
|
||||
```
|
||||
|
||||
We've also prepared preprocessed binary formatted files, which can be found in the container below:
|
||||
|
||||
```bash
|
||||
root@iZj6caifjouj5yu8xgsiysZ:/home/zvec# ls -al /tmp/cohere/*zvec
|
||||
/tmp/cohere/cohere_large_10m_zvec:
|
||||
total 30204572
|
||||
drwxr-xr-x 2 root root 4096 Feb 5 13:12 .
|
||||
drwxr-xr-x 6 root root 4096 Feb 6 03:38 ..
|
||||
-rw-r--r-- 1 root root 8664837 Feb 5 13:06 cohere_test_vector_10m.1000.new.txt
|
||||
-rw-r--r-- 1 root root 30920004295 Feb 5 13:04 cohere_train_vector_10m.new.zvec.vecs
|
||||
-rw-r--r-- 1 root root 792835 Feb 5 13:05 neighbors.txt
|
||||
|
||||
/tmp/cohere/cohere_medium_1m_zvec:
|
||||
total 3028688
|
||||
drwxr-xr-x 2 root root 4096 Feb 5 13:14 .
|
||||
drwxr-xr-x 6 root root 4096 Feb 6 03:38 ..
|
||||
-rw-r--r-- 1 root root 8661108 Feb 5 13:07 cohere_test_vector_1m.1000.new.txt
|
||||
-rw-r--r-- 1 root root 3092004295 Feb 5 13:08 cohere_train_vector_1m.new.zvec.vecs
|
||||
-rw-r--r-- 1 root root 692969 Feb 5 13:08 neighbors.txt
|
||||
```
|
||||
|
||||
### Preparing Bench Config
|
||||
Prepare Build Config
|
||||
|
||||
```yaml
|
||||
BuilderCommon:
|
||||
BuilderClass: HnswStreamer
|
||||
BuildFile: /tmp/cohere/cohere_large_10m_zvec/cohere_train_vector_10m.zvec.vecs
|
||||
NeedTrain: true
|
||||
TrainFile: /tmp/cohere/cohere_large_10m_zvec/cohere_train_vector_10m.zvec.vecs
|
||||
DumpPath: /home/zvec/bench/config/cohere_train_vector_10m.dump.index
|
||||
IndexPath: /home/zvec/bench/config/cohere_train_vector_10m.index
|
||||
|
||||
ConverterName: CosineInt8Converter
|
||||
MetricName: Cosine
|
||||
|
||||
ThreadCount: 16
|
||||
|
||||
BuilderParams:
|
||||
proxima.general.builder.thread_count: !!int 16
|
||||
proxima.hnsw.builder.thread_count: !!int 16
|
||||
```
|
||||
|
||||
Prepare Search Config
|
||||
|
||||
```yaml
|
||||
SearcherCommon:
|
||||
SearcherClass: HnswStreamer
|
||||
IndexPath: /home/zvec/bench/config/cohere_train_vector_10m.index
|
||||
TopK: 1,10,50,100
|
||||
QueryFile: /tmp/cohere/cohere_large_10m_zvec/cohere_test_vector_1000.new.txt
|
||||
QueryType: float
|
||||
QueryFirstSep: ";"
|
||||
QuerySecondSep: " "
|
||||
GroundTruthFile: /tmp/cohere/cohere_large_10m_zvec/neighbors.txt
|
||||
RecallThreadCount: 1
|
||||
BenchThreadCount: 16
|
||||
BenchIterCount: 1000000000
|
||||
CompareById: true
|
||||
|
||||
SearcherParams:
|
||||
proxima.hnsw.streamer.ef: !!int 250
|
||||
```
|
||||
|
||||
### Building Index
|
||||
Conduct Build
|
||||
```bash
|
||||
$ /home/zvec/workspace/zvec/build/bin/local_build_original ./build.yaml
|
||||
```
|
||||
|
||||
### Performing Bench
|
||||
Conduct Recall
|
||||
```bash
|
||||
$ /home/zvec/workspace/zvec/build/bin/recall_original ./search.yaml
|
||||
```
|
||||
|
||||
Conduct Bench
|
||||
```bash
|
||||
$ /home/zvec/workspace/zvec/build/bin/bench_original ./search.yaml
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,911 @@
|
||||
// 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 "bench_result.h"
|
||||
#include "helper.h"
|
||||
|
||||
static bool g_debug_mode = 0;
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Bench
|
||||
//------------------------------------------------------------
|
||||
enum RetrievalMode { RM_UNDEFINED = 0, RM_DENSE = 1, RM_SPARSE = 2 };
|
||||
|
||||
enum FilterMode { FM_UNDEFINED = 0, FM_NONE = 1, FM_TAG = 2 };
|
||||
|
||||
template <typename T>
|
||||
class Bench {
|
||||
public:
|
||||
Bench(size_t threads, size_t bench_secs, size_t batch_count,
|
||||
RetrievalMode &retrieval_mode, FilterMode filter_mode)
|
||||
: threads_(threads),
|
||||
bench_secs_(bench_secs),
|
||||
batch_count_(batch_count),
|
||||
retrieval_mode_{retrieval_mode},
|
||||
filter_mode_{filter_mode} {
|
||||
if (threads_ == 0) {
|
||||
pool_ = make_shared<ThreadPool>(false);
|
||||
threads_ = pool_->count();
|
||||
cout << "Using cpu count as thread pool count[" << threads_ << "]"
|
||||
<< endl;
|
||||
} else {
|
||||
pool_ = make_shared<ThreadPool>(threads_, false);
|
||||
cout << "Using thread pool count[" << threads_ << "]" << endl;
|
||||
}
|
||||
if (batch_count_ < 1) {
|
||||
batch_count_ = 1;
|
||||
}
|
||||
}
|
||||
|
||||
static void stop(int signo) {
|
||||
if (STOP_NOW) {
|
||||
exit(signo);
|
||||
}
|
||||
STOP_NOW = true;
|
||||
cout << "\rTrying to stop. press [Ctrl+C] again kill immediately." << endl
|
||||
<< flush;
|
||||
}
|
||||
|
||||
bool load_query(const std::string &query_file, const std::string &first_sep,
|
||||
const std::string &second_sep) {
|
||||
TxtInputReader<T> reader;
|
||||
vector<vector<T>> queries;
|
||||
vector<SparseData<T>> sparse_data;
|
||||
vector<vector<uint64_t>> taglists;
|
||||
|
||||
if (!reader.load_query(query_file, first_sep, second_sep, queries,
|
||||
sparse_data, taglists)) {
|
||||
LOG_ERROR("Load query error");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (batch_count_ == 1) {
|
||||
batch_queries_ = queries;
|
||||
|
||||
for (size_t i = 0; i < sparse_data.size(); ++i) {
|
||||
vector<uint32_t> sparse_count;
|
||||
sparse_count.push_back(sparse_data[i].count);
|
||||
|
||||
batch_sparse_counts_.push_back(sparse_count);
|
||||
batch_sparse_indices_.push_back(sparse_data[i].indices);
|
||||
batch_sparse_features_.push_back(sparse_data[i].features);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < taglists.size(); ++i) {
|
||||
vector<vector<uint64_t>> new_taglists;
|
||||
new_taglists.push_back(taglists[i]);
|
||||
|
||||
batch_taglists_.push_back(std::move(new_taglists));
|
||||
}
|
||||
} else {
|
||||
size_t num_batch = (queries.size() + batch_count_ - 1) / batch_count_;
|
||||
size_t idx = 0;
|
||||
for (size_t n = 0; n < num_batch; ++n) {
|
||||
vector<T> batch_query;
|
||||
vector<uint32_t> batch_sparse_count;
|
||||
vector<uint32_t> batch_sparse_indices;
|
||||
vector<T> batch_sparse_feature;
|
||||
vector<vector<uint64_t>> batch_taglists;
|
||||
|
||||
for (size_t i = 0; i < batch_count_; ++i) {
|
||||
for (size_t k = 0; k < queries[idx].size(); ++k) {
|
||||
batch_query.push_back(queries[idx][k]);
|
||||
}
|
||||
|
||||
batch_sparse_count.push_back(sparse_data[idx].count);
|
||||
|
||||
for (size_t k = 0; k < sparse_data[idx].indices.size(); ++k) {
|
||||
batch_sparse_indices.push_back(sparse_data[idx].indices[k]);
|
||||
}
|
||||
|
||||
for (size_t k = 0; k < sparse_data[idx].features.size(); ++k) {
|
||||
batch_sparse_feature.push_back(sparse_data[idx].features[k]);
|
||||
}
|
||||
|
||||
if (taglists.size() > idx) {
|
||||
batch_taglists.push_back(taglists[idx]);
|
||||
}
|
||||
|
||||
idx = (idx + 1) % queries.size();
|
||||
}
|
||||
|
||||
batch_queries_.push_back(batch_query);
|
||||
batch_sparse_counts_.push_back(batch_sparse_count);
|
||||
batch_sparse_indices_.push_back(batch_sparse_indices);
|
||||
batch_sparse_features_.push_back(batch_sparse_feature);
|
||||
batch_taglists_.push_back(batch_taglists);
|
||||
}
|
||||
}
|
||||
|
||||
dim_ = queries[0].size();
|
||||
if (typeid(T) == typeid(float)) {
|
||||
qmeta_.set_meta(IndexMeta::DataType::DT_FP32, dim_);
|
||||
} else if (typeid(T) == typeid(int8_t)) {
|
||||
qmeta_.set_meta(IndexMeta::DataType::DT_INT8, dim_);
|
||||
} else {
|
||||
LOG_ERROR("unsupported type");
|
||||
return false;
|
||||
}
|
||||
|
||||
cout << "Load query done!" << endl;
|
||||
return true;
|
||||
}
|
||||
|
||||
void run(core_interface::Index::Pointer index,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param,
|
||||
int max_iter, int topk) {
|
||||
// Check
|
||||
if (batch_queries_.size() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
query_param_ = query_param;
|
||||
query_param_->topk = topk;
|
||||
query_param_->is_linear = false;
|
||||
|
||||
// Do bench
|
||||
signal(SIGINT, stop);
|
||||
bench_result_.mark_start();
|
||||
auto start_time = Monotime::MilliSeconds();
|
||||
for (size_t i = 0; i < threads_; ++i) {
|
||||
pool_->execute(this, &Bench<T>::start_bench, index, max_iter, &STOP_NOW);
|
||||
}
|
||||
|
||||
while (!pool_->is_finished()) {
|
||||
this_thread::sleep_for(chrono::milliseconds(1));
|
||||
if (Monotime::MilliSeconds() - start_time > bench_secs_ * 1000) {
|
||||
STOP_NOW = true;
|
||||
}
|
||||
}
|
||||
|
||||
pool_->wait_finish();
|
||||
|
||||
bench_result_.mark_end();
|
||||
bench_result_.print();
|
||||
}
|
||||
|
||||
void set_tag_lists(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
|
||||
const std::vector<uint64_t> &tag_key_list) {
|
||||
id_to_tags_list_ = id_to_tags_list;
|
||||
tag_key_list_ = tag_key_list;
|
||||
}
|
||||
|
||||
private:
|
||||
void start_bench(core_interface::Index::Pointer index, size_t max_iter,
|
||||
const bool *is_stop) {
|
||||
size_t thread_index = pool_->indexof_this();
|
||||
|
||||
size_t i = thread_index;
|
||||
for (; i < max_iter && !*is_stop; i += threads_) {
|
||||
int idx = i % batch_queries_.size();
|
||||
|
||||
// prefilter
|
||||
FilterResultCache filter_cache;
|
||||
std::shared_ptr<IndexFilter> filter_ptr = nullptr;
|
||||
if (filter_mode_ == FM_TAG) {
|
||||
if (batch_taglists_[idx].size() != 1) {
|
||||
LOG_ERROR("query tag list not equal to one!");
|
||||
return;
|
||||
}
|
||||
|
||||
int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[idx][0],
|
||||
tag_key_list_);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("prefilter failed, idx: %d", idx);
|
||||
return;
|
||||
}
|
||||
|
||||
auto filterFunc = [&](uint64_t key) { return filter_cache.find(key); };
|
||||
|
||||
filter_ptr = std::make_shared<IndexFilter>();
|
||||
filter_ptr->set(filterFunc);
|
||||
}
|
||||
|
||||
auto query_param = query_param_->Clone();
|
||||
query_param->filter = filter_ptr;
|
||||
|
||||
|
||||
// Do knn_search
|
||||
uint64_t start = Monotime::MicroSeconds();
|
||||
int ret;
|
||||
if (retrieval_mode_ == RM_DENSE) {
|
||||
if (batch_count_ == 1) {
|
||||
ret = do_knn_search<T>(index, batch_queries_[idx], query_param);
|
||||
} else {
|
||||
ret = do_knn_search_batch<T>(index, batch_queries_[idx], query_param);
|
||||
}
|
||||
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Failed to knn search, ret=%d %s", ret,
|
||||
IndexError::What(ret));
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
std::string mode = retrieval_mode_ == 1 ? "Dense" : "Sparse";
|
||||
LOG_ERROR("unsupported retrieval mode: %s", mode.c_str());
|
||||
}
|
||||
|
||||
uint64_t end = Monotime::MicroSeconds();
|
||||
|
||||
// Do sample
|
||||
bench_result_.add_time(batch_count_, end - start);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
typename std::enable_if<
|
||||
std::is_same<float, U>::value || std::is_same<int8_t, U>::value ||
|
||||
std::is_same<uint32_t, U>::value || std::is_same<uint64_t, U>::value,
|
||||
int>::type
|
||||
do_knn_search(core_interface::Index::Pointer index, const vector<U> &query,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param) {
|
||||
core_interface::DenseVector dense_query;
|
||||
dense_query.data = query.data();
|
||||
core_interface::VectorData query_data;
|
||||
query_data.vector = dense_query;
|
||||
|
||||
core_interface::SearchResult search_result;
|
||||
int ret = index->Search(query_data, query_param, &search_result);
|
||||
if (ret < 0) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
if (search_result.doc_list_.empty()) {
|
||||
LOG_ERROR("Search results is empty");
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
typename std::enable_if<
|
||||
std::is_same<float, U>::value || std::is_same<int8_t, U>::value ||
|
||||
std::is_same<uint32_t, U>::value || std::is_same<uint64_t, U>::value,
|
||||
int>::type
|
||||
do_knn_search_batch(
|
||||
core_interface::Index::Pointer index, const vector<U> &query,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param) {
|
||||
// For batch search, we search each query separately
|
||||
size_t qnum = query.size() / dim_;
|
||||
for (size_t i = 0; i < qnum; ++i) {
|
||||
core_interface::DenseVector dense_query;
|
||||
dense_query.data = query.data() + i * dim_;
|
||||
core_interface::VectorData query_data;
|
||||
query_data.vector = dense_query;
|
||||
|
||||
core_interface::SearchResult search_result;
|
||||
int ret = index->Search(query_data, query_param, &search_result);
|
||||
if (ret < 0) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
if (search_result.doc_list_.empty()) {
|
||||
LOG_ERROR("Search results is empty for batch query %zu", i);
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
private:
|
||||
IndexQueryMeta qmeta_{};
|
||||
size_t threads_;
|
||||
size_t bench_secs_;
|
||||
size_t batch_count_;
|
||||
size_t dim_;
|
||||
shared_ptr<ThreadPool> pool_;
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param_;
|
||||
|
||||
vector<vector<T>> batch_queries_;
|
||||
vector<vector<uint32_t>> batch_sparse_counts_;
|
||||
vector<vector<uint32_t>> batch_sparse_indices_;
|
||||
vector<vector<T>> batch_sparse_features_;
|
||||
vector<vector<vector<uint64_t>>> batch_taglists_;
|
||||
|
||||
// Tag lists for filtering
|
||||
std::vector<std::vector<uint64_t>> id_to_tags_list_;
|
||||
std::vector<uint64_t> tag_key_list_;
|
||||
|
||||
BenchResult bench_result_;
|
||||
RetrievalMode retrieval_mode_{RM_UNDEFINED};
|
||||
FilterMode filter_mode_{FM_NONE};
|
||||
static bool STOP_NOW;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
bool Bench<T>::STOP_NOW = false;
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Sparse Bench
|
||||
//------------------------------------------------------------
|
||||
template <typename T>
|
||||
class SparseBench {
|
||||
public:
|
||||
SparseBench(size_t threads, size_t bench_secs, size_t batch_count,
|
||||
FilterMode filter_mode)
|
||||
: threads_(threads),
|
||||
bench_secs_(bench_secs),
|
||||
batch_count_(batch_count),
|
||||
filter_mode_{filter_mode} {
|
||||
if (threads_ == 0) {
|
||||
pool_ = make_shared<ThreadPool>(false);
|
||||
threads_ = pool_->count();
|
||||
cout << "Using cpu count as thread pool count[" << threads_ << "]"
|
||||
<< endl;
|
||||
} else {
|
||||
pool_ = make_shared<ThreadPool>(threads_, false);
|
||||
cout << "Using thread pool count[" << threads_ << "]" << endl;
|
||||
}
|
||||
if (batch_count_ < 1) {
|
||||
batch_count_ = 1;
|
||||
}
|
||||
}
|
||||
|
||||
static void stop(int signo) {
|
||||
if (STOP_NOW) {
|
||||
exit(signo);
|
||||
}
|
||||
STOP_NOW = true;
|
||||
cout << "\rTrying to stop. press [Ctrl+C] again kill immediately." << endl
|
||||
<< flush;
|
||||
}
|
||||
|
||||
bool load_query(const std::string &query_file, const std::string &first_sep,
|
||||
const std::string &second_sep) {
|
||||
TxtInputReader<T> reader;
|
||||
vector<vector<T>> queries;
|
||||
vector<SparseData<T>> sparse_data;
|
||||
vector<vector<uint64_t>> taglists;
|
||||
|
||||
if (!reader.load_query(query_file, first_sep, second_sep, queries,
|
||||
sparse_data, taglists)) {
|
||||
LOG_ERROR("Load query error");
|
||||
return false;
|
||||
}
|
||||
|
||||
linear_sparse_data_ = sparse_data;
|
||||
|
||||
if (batch_count_ == 1) {
|
||||
for (size_t i = 0; i < sparse_data.size(); ++i) {
|
||||
vector<uint32_t> sparse_count;
|
||||
sparse_count.push_back(sparse_data[i].count);
|
||||
|
||||
batch_sparse_counts_.push_back(sparse_count);
|
||||
batch_sparse_indices_.push_back(sparse_data[i].indices);
|
||||
batch_sparse_features_.push_back(sparse_data[i].features);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < taglists.size(); ++i) {
|
||||
vector<vector<uint64_t>> new_taglists;
|
||||
new_taglists.push_back(taglists[i]);
|
||||
|
||||
batch_taglists_.push_back(std::move(new_taglists));
|
||||
}
|
||||
} else {
|
||||
size_t num_batch = (queries.size() + batch_count_ - 1) / batch_count_;
|
||||
size_t idx = 0;
|
||||
for (size_t n = 0; n < num_batch; ++n) {
|
||||
vector<uint32_t> batch_sparse_count;
|
||||
vector<uint32_t> batch_sparse_indices;
|
||||
vector<T> batch_sparse_feature;
|
||||
vector<vector<uint64_t>> batch_taglists;
|
||||
|
||||
for (size_t i = 0; i < batch_count_; ++i) {
|
||||
batch_sparse_count.push_back(sparse_data[idx].count);
|
||||
|
||||
for (size_t k = 0; k < sparse_data[idx].indices.size(); ++k) {
|
||||
batch_sparse_indices.push_back(sparse_data[idx].indices[k]);
|
||||
}
|
||||
|
||||
for (size_t k = 0; k < sparse_data[idx].features.size(); ++k) {
|
||||
batch_sparse_feature.push_back(sparse_data[idx].features[k]);
|
||||
}
|
||||
|
||||
if (taglists.size() > idx) {
|
||||
batch_taglists.push_back(taglists[idx]);
|
||||
}
|
||||
|
||||
idx = (idx + 1) % queries.size();
|
||||
}
|
||||
|
||||
batch_sparse_counts_.push_back(batch_sparse_count);
|
||||
batch_sparse_indices_.push_back(batch_sparse_indices);
|
||||
batch_sparse_features_.push_back(batch_sparse_feature);
|
||||
batch_taglists_.push_back(batch_taglists);
|
||||
}
|
||||
}
|
||||
|
||||
if (typeid(T) == typeid(float)) {
|
||||
qmeta_.set_data_type(IndexMeta::DataType::DT_FP32);
|
||||
} else if (typeid(T) == typeid(int8_t)) {
|
||||
qmeta_.set_data_type(IndexMeta::DataType::DT_INT8);
|
||||
} else {
|
||||
LOG_ERROR("unsupported type");
|
||||
return false;
|
||||
}
|
||||
|
||||
cout << "Load query done!" << endl;
|
||||
return true;
|
||||
}
|
||||
|
||||
void run(core_interface::Index::Pointer index,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param,
|
||||
int max_iter, int topk) {
|
||||
// Check
|
||||
if (batch_sparse_counts_.size() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
query_param_ = query_param;
|
||||
query_param_->topk = topk;
|
||||
query_param_->is_linear = false;
|
||||
|
||||
// Do bench
|
||||
signal(SIGINT, stop);
|
||||
bench_result_.mark_start();
|
||||
auto start_time = Monotime::MilliSeconds();
|
||||
for (size_t i = 0; i < threads_; ++i) {
|
||||
pool_->execute(this, &SparseBench<T>::start_bench, index, max_iter,
|
||||
&STOP_NOW);
|
||||
}
|
||||
|
||||
while (!pool_->is_finished()) {
|
||||
this_thread::sleep_for(chrono::milliseconds(1));
|
||||
if (Monotime::MilliSeconds() - start_time > bench_secs_ * 1000) {
|
||||
STOP_NOW = true;
|
||||
}
|
||||
}
|
||||
|
||||
pool_->wait_finish();
|
||||
|
||||
bench_result_.mark_end();
|
||||
bench_result_.print();
|
||||
}
|
||||
|
||||
void set_tag_lists(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
|
||||
const std::vector<uint64_t> &tag_key_list) {
|
||||
id_to_tags_list_ = id_to_tags_list;
|
||||
tag_key_list_ = tag_key_list;
|
||||
}
|
||||
|
||||
private:
|
||||
void start_bench(core_interface::Index::Pointer index, size_t max_iter,
|
||||
const bool *is_stop) {
|
||||
size_t thread_index = pool_->indexof_this();
|
||||
|
||||
size_t i = thread_index;
|
||||
size_t sparse_query_size = batch_sparse_indices_.size();
|
||||
for (; i < max_iter && !*is_stop; i += threads_) {
|
||||
int idx = i % sparse_query_size;
|
||||
|
||||
// prefilter
|
||||
FilterResultCache filter_cache;
|
||||
std::shared_ptr<IndexFilter> filter_ptr = nullptr;
|
||||
if (filter_mode_ == FM_TAG) {
|
||||
if (batch_taglists_[idx].size() != 1) {
|
||||
LOG_ERROR("query tag list not equal to one!");
|
||||
return;
|
||||
}
|
||||
|
||||
int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[idx][0],
|
||||
tag_key_list_);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("prefilter failed, idx: %d", idx);
|
||||
return;
|
||||
}
|
||||
|
||||
auto filterFunc = [&](uint64_t key) { return filter_cache.find(key); };
|
||||
|
||||
filter_ptr = std::make_shared<IndexFilter>();
|
||||
filter_ptr->set(filterFunc);
|
||||
}
|
||||
|
||||
auto query_param = query_param_->Clone();
|
||||
query_param->filter = filter_ptr;
|
||||
|
||||
// Do knn_search
|
||||
uint64_t start = Monotime::MicroSeconds();
|
||||
int ret;
|
||||
if (batch_count_ == 1) {
|
||||
if (batch_sparse_counts_[idx].size() != 1) {
|
||||
LOG_ERROR("Sparse count size should be 1, since batch count is 1");
|
||||
return;
|
||||
}
|
||||
ret = do_knn_search<T>(index, batch_sparse_counts_[idx][0],
|
||||
batch_sparse_indices_[idx],
|
||||
batch_sparse_features_[idx], query_param);
|
||||
} else {
|
||||
ret = do_knn_search_batch<T>(
|
||||
index, batch_sparse_counts_[idx], batch_sparse_indices_[idx],
|
||||
batch_sparse_features_[idx], idx, query_param);
|
||||
}
|
||||
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Failed to sparse knn search, ret=%d %s", ret,
|
||||
IndexError::What(ret));
|
||||
return;
|
||||
}
|
||||
|
||||
uint64_t end = Monotime::MicroSeconds();
|
||||
|
||||
// Do sample
|
||||
bench_result_.add_time(batch_count_, end - start);
|
||||
}
|
||||
}
|
||||
|
||||
// sparse search - single query
|
||||
template <typename U>
|
||||
typename std::enable_if<std::is_same<float, U>::value, int>::type
|
||||
do_knn_search(core_interface::Index::Pointer index,
|
||||
const uint32_t sparse_count,
|
||||
const vector<uint32_t> &sparse_indices,
|
||||
const vector<U> &sparse_feature,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param) {
|
||||
core_interface::SparseVector sparse_query;
|
||||
sparse_query.count = sparse_count;
|
||||
sparse_query.indices = sparse_indices.data();
|
||||
sparse_query.values = sparse_feature.data();
|
||||
core_interface::VectorData query_data;
|
||||
query_data.vector = sparse_query;
|
||||
|
||||
core_interface::SearchResult search_result;
|
||||
int ret = index->Search(query_data, query_param, &search_result);
|
||||
if (ret < 0) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
if (search_result.doc_list_.empty()) {
|
||||
LOG_ERROR("Search results is empty");
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
typename std::enable_if<std::is_same<int8_t, U>::value ||
|
||||
std::is_same<uint32_t, U>::value ||
|
||||
std::is_same<uint64_t, U>::value,
|
||||
int>::type
|
||||
do_knn_search(core_interface::Index::Pointer /*index*/,
|
||||
const uint32_t /*sparse_count*/,
|
||||
const vector<uint32_t> & /*sparse_indices*/,
|
||||
const vector<U> & /*sparse_feature*/,
|
||||
core_interface::BaseIndexQueryParam::Pointer /*query_param*/) {
|
||||
return IndexError_Unsupported;
|
||||
}
|
||||
|
||||
// sparse search - batch
|
||||
template <typename U>
|
||||
typename std::enable_if<std::is_same<float, U>::value, int>::type
|
||||
do_knn_search_batch(
|
||||
core_interface::Index::Pointer index,
|
||||
const vector<uint32_t> &sparse_count,
|
||||
const vector<uint32_t> & /*sparse_indices*/,
|
||||
const vector<U> & /*sparse_feature*/, size_t batch_idx,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param) {
|
||||
// For batch search, search each query separately
|
||||
for (size_t i = 0; i < sparse_count.size(); ++i) {
|
||||
size_t query_idx = batch_idx * batch_count_ + i;
|
||||
if (query_idx >= linear_sparse_data_.size()) {
|
||||
break;
|
||||
}
|
||||
|
||||
const auto &single_sparse = linear_sparse_data_[query_idx];
|
||||
core_interface::SparseVector sparse_query;
|
||||
sparse_query.count = single_sparse.count;
|
||||
sparse_query.indices = single_sparse.indices.data();
|
||||
sparse_query.values = single_sparse.features.data();
|
||||
core_interface::VectorData query_data;
|
||||
query_data.vector = sparse_query;
|
||||
|
||||
core_interface::SearchResult search_result;
|
||||
int ret = index->Search(query_data, query_param, &search_result);
|
||||
if (ret < 0) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
if (search_result.doc_list_.empty()) {
|
||||
LOG_ERROR("Search results is empty for batch query %zu", i);
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
typename std::enable_if<std::is_same<int8_t, U>::value ||
|
||||
std::is_same<uint32_t, U>::value ||
|
||||
std::is_same<uint64_t, U>::value,
|
||||
int>::type
|
||||
do_knn_search_batch(
|
||||
core_interface::Index::Pointer /*index*/,
|
||||
const vector<uint32_t> & /*sparse_count*/,
|
||||
const vector<uint32_t> & /*sparse_indices*/,
|
||||
const vector<U> & /*sparse_feature*/, size_t /*batch_idx*/,
|
||||
core_interface::BaseIndexQueryParam::Pointer /*query_param*/) {
|
||||
return IndexError_Unsupported;
|
||||
}
|
||||
|
||||
private:
|
||||
IndexQueryMeta qmeta_{};
|
||||
size_t threads_;
|
||||
size_t bench_secs_;
|
||||
size_t batch_count_;
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param_;
|
||||
shared_ptr<ThreadPool> pool_;
|
||||
|
||||
vector<SparseData<T>> linear_sparse_data_;
|
||||
vector<vector<uint32_t>> batch_sparse_counts_;
|
||||
vector<vector<uint32_t>> batch_sparse_indices_;
|
||||
vector<vector<T>> batch_sparse_features_;
|
||||
vector<vector<vector<uint64_t>>> batch_taglists_;
|
||||
|
||||
// Tag lists for filtering
|
||||
std::vector<std::vector<uint64_t>> id_to_tags_list_;
|
||||
std::vector<uint64_t> tag_key_list_;
|
||||
|
||||
FilterMode filter_mode_{FM_NONE};
|
||||
BenchResult bench_result_;
|
||||
static bool STOP_NOW;
|
||||
};
|
||||
template <typename T>
|
||||
bool SparseBench<T>::STOP_NOW = false;
|
||||
|
||||
bool check_config(YAML::Node &config_node) {
|
||||
auto common = config_node["IndexCommon"];
|
||||
if (!common) {
|
||||
LOG_ERROR("Can not find [IndexCommon] in config");
|
||||
return false;
|
||||
}
|
||||
if (!common["IndexConfig"]) {
|
||||
LOG_ERROR("Can not find [IndexConfig] in config");
|
||||
return false;
|
||||
}
|
||||
if (!common["IndexPath"]) {
|
||||
LOG_ERROR("Can not find [IndexPath] in config");
|
||||
return false;
|
||||
}
|
||||
if (!common["TopK"]) {
|
||||
LOG_ERROR("Can not find [TopK] in config");
|
||||
return false;
|
||||
}
|
||||
if (!common["QueryFile"]) {
|
||||
LOG_ERROR("Can not find [QueryFile] in config");
|
||||
return false;
|
||||
}
|
||||
|
||||
auto query_config = config_node["QueryConfig"];
|
||||
if (!query_config) {
|
||||
LOG_ERROR("Can not find [QueryConfig] in config");
|
||||
return false;
|
||||
}
|
||||
if (!query_config["QueryParam"]) {
|
||||
LOG_ERROR("Can not find [QueryConfig.QueryParam] in config");
|
||||
return false;
|
||||
}
|
||||
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
void usage(void) {
|
||||
cout << "Usage: bench CONFIG.yaml [plugin file path]" << endl;
|
||||
}
|
||||
|
||||
|
||||
int bench(std::string &query_type, size_t thread_count, size_t batch_count,
|
||||
size_t top_k, string query_file, string &first_sep,
|
||||
string &second_sep, size_t bench_secs, size_t iter_count,
|
||||
core_interface::Index::Pointer index,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param,
|
||||
string &index_dir, RetrievalMode retrieval_mode,
|
||||
FilterMode filter_mode) {
|
||||
if (filter_mode == FM_TAG && batch_count > 1) {
|
||||
LOG_ERROR("filter mode can not be run in batch mode");
|
||||
return -1;
|
||||
}
|
||||
|
||||
std::vector<std::vector<uint64_t>> id_to_tags_list;
|
||||
std::vector<uint64_t> tag_key_list;
|
||||
// Load tag lists if available
|
||||
load_taglists(index_dir, id_to_tags_list, tag_key_list);
|
||||
|
||||
if (query_type == "float") {
|
||||
Bench<float> bench(thread_count, bench_secs, batch_count, retrieval_mode,
|
||||
filter_mode);
|
||||
bench.load_query(query_file, first_sep, second_sep);
|
||||
bench.set_tag_lists(id_to_tags_list, tag_key_list);
|
||||
bench.run(index, query_param, iter_count, top_k);
|
||||
} else if (query_type == "int8") {
|
||||
Bench<int8_t> bench(thread_count, bench_secs, batch_count, retrieval_mode,
|
||||
filter_mode);
|
||||
bench.load_query(query_file, first_sep, second_sep);
|
||||
bench.set_tag_lists(id_to_tags_list, tag_key_list);
|
||||
bench.run(index, query_param, iter_count, top_k);
|
||||
} else if (query_type == "binary") {
|
||||
Bench<uint32_t> bench(thread_count, bench_secs, batch_count, retrieval_mode,
|
||||
filter_mode);
|
||||
bench.load_query(query_file, first_sep, second_sep);
|
||||
bench.set_tag_lists(id_to_tags_list, tag_key_list);
|
||||
bench.run(index, query_param, iter_count, top_k);
|
||||
} else if (query_type == "binary64") {
|
||||
Bench<uint64_t> bench(thread_count, bench_secs, batch_count, retrieval_mode,
|
||||
filter_mode);
|
||||
bench.load_query(query_file, first_sep, second_sep);
|
||||
bench.set_tag_lists(id_to_tags_list, tag_key_list);
|
||||
bench.run(index, query_param, iter_count, top_k);
|
||||
} else {
|
||||
LOG_ERROR("Can not recognize type: %s", query_type.c_str());
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int bench_sparse(std::string &query_type, size_t thread_count,
|
||||
size_t batch_count, size_t top_k, string query_file,
|
||||
string &first_sep, string &second_sep, size_t bench_secs,
|
||||
size_t iter_count, core_interface::Index::Pointer index,
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param,
|
||||
string &index_dir, FilterMode filter_mode) {
|
||||
if (filter_mode == FM_TAG && batch_count > 1) {
|
||||
LOG_ERROR("filter mode can not be run in batch mode");
|
||||
return -1;
|
||||
}
|
||||
|
||||
std::vector<std::vector<uint64_t>> id_to_tags_list;
|
||||
std::vector<uint64_t> tag_key_list;
|
||||
// Load tag lists if available
|
||||
load_taglists(index_dir, id_to_tags_list, tag_key_list);
|
||||
|
||||
if (query_type == "float") {
|
||||
SparseBench<float> bench(thread_count, bench_secs, batch_count,
|
||||
filter_mode);
|
||||
bench.load_query(query_file, first_sep, second_sep);
|
||||
bench.set_tag_lists(id_to_tags_list, tag_key_list);
|
||||
bench.run(index, query_param, iter_count, top_k);
|
||||
} else if (query_type == "int8") {
|
||||
SparseBench<int8_t> bench(thread_count, bench_secs, batch_count,
|
||||
filter_mode);
|
||||
bench.load_query(query_file, first_sep, second_sep);
|
||||
bench.set_tag_lists(id_to_tags_list, tag_key_list);
|
||||
bench.run(index, query_param, iter_count, top_k);
|
||||
} else {
|
||||
LOG_ERROR("Can not recognize type: %s", query_type.c_str());
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int main(int argc, char *argv[]) {
|
||||
if (argc < 2) {
|
||||
usage();
|
||||
return -1;
|
||||
}
|
||||
|
||||
IndexPluginBroker broker;
|
||||
std::string error;
|
||||
for (int i = 2; i < argc; ++i) {
|
||||
if (!broker.emplace(argv[i], &error)) {
|
||||
LOG_ERROR("Failed to load plugin: %s (%s)", argv[i], error.c_str());
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
YAML::Node config_node;
|
||||
try {
|
||||
config_node = YAML::LoadFile(argv[1]);
|
||||
} catch (...) {
|
||||
LOG_ERROR("Load YAML file[%s] failed!", argv[1]);
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (!check_config(config_node)) {
|
||||
return -1;
|
||||
}
|
||||
auto config_common = config_node["IndexCommon"];
|
||||
|
||||
map<string, int> LOG_LEVEL = {{"debug", IndexLogger::LEVEL_DEBUG},
|
||||
{"info", IndexLogger::LEVEL_INFO},
|
||||
{"warn", IndexLogger::LEVEL_WARN},
|
||||
{"error", IndexLogger::LEVEL_ERROR},
|
||||
{"fatal", IndexLogger::LEVEL_FATAL}};
|
||||
string log_level = config_common["LogLevel"]
|
||||
? config_common["LogLevel"].as<string>()
|
||||
: "debug";
|
||||
transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower);
|
||||
if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) {
|
||||
IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]);
|
||||
zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]);
|
||||
}
|
||||
|
||||
// Calculate Bench
|
||||
size_t thread_count = config_common["BenchThreadCount"]
|
||||
? config_common["BenchThreadCount"].as<uint64_t>()
|
||||
: 0;
|
||||
size_t iter_count = config_common["BenchIterCount"]
|
||||
? config_common["BenchIterCount"].as<uint64_t>()
|
||||
: 10000;
|
||||
size_t batch_count = config_common["BenchBatchCount"]
|
||||
? config_common["BenchBatchCount"].as<uint64_t>()
|
||||
: 0;
|
||||
g_debug_mode = config_common["DebugMode"]
|
||||
? config_common["DebugMode"].as<bool>()
|
||||
: false;
|
||||
string topk_str = config_common["TopK"].as<string>();
|
||||
|
||||
RetrievalMode retrieval_mode{RM_DENSE};
|
||||
if (config_common["RetrievalMode"]) {
|
||||
std::string retrieval_mode_str =
|
||||
config_common["RetrievalMode"].as<string>();
|
||||
if (retrieval_mode_str == "dense") {
|
||||
retrieval_mode = RM_DENSE;
|
||||
} else if (retrieval_mode_str == "sparse") {
|
||||
retrieval_mode = RM_SPARSE;
|
||||
}
|
||||
}
|
||||
|
||||
FilterMode filter_mode{FM_NONE};
|
||||
if (config_common["FilterMode"]) {
|
||||
std::string filter_mode_str = config_common["FilterMode"].as<string>();
|
||||
if (filter_mode_str == "tag") {
|
||||
filter_mode = FM_TAG;
|
||||
}
|
||||
}
|
||||
|
||||
vector<int32_t> topk_values;
|
||||
StringHelper::Split(topk_str, ",", &topk_values);
|
||||
size_t top_k = *topk_values.rbegin();
|
||||
string query_file = config_common["QueryFile"].as<string>();
|
||||
string first_sep = config_common["QueryFirstSep"]
|
||||
? config_common["QueryFirstSep"].as<string>()
|
||||
: ";";
|
||||
string second_sep = config_common["QuerySecondSep"]
|
||||
? config_common["QuerySecondSep"].as<string>()
|
||||
: " ";
|
||||
string query_type = config_common["QueryType"]
|
||||
? config_common["QueryType"].as<string>()
|
||||
: "float";
|
||||
size_t bench_secs = config_common["BenchSecs"]
|
||||
? config_common["BenchSecs"].as<uint64_t>()
|
||||
: 60;
|
||||
|
||||
string index_dir = config_common["IndexPath"].as<string>();
|
||||
|
||||
core_interface::Index::Pointer index;
|
||||
core_interface::BaseIndexQueryParam::Pointer query_param;
|
||||
if (0 !=
|
||||
parse_and_load_index_param(config_node, index_dir, index, query_param)) {
|
||||
LOG_ERROR("Failed to parse and load index param");
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (retrieval_mode == RM_SPARSE) {
|
||||
bench_sparse(query_type, thread_count, batch_count, top_k, query_file,
|
||||
first_sep, second_sep, bench_secs, iter_count, index,
|
||||
query_param, index_dir, filter_mode);
|
||||
|
||||
cout << "Bench Sparse done." << endl;
|
||||
} else {
|
||||
bench(query_type, thread_count, batch_count, top_k, query_file, first_sep,
|
||||
second_sep, bench_secs, iter_count, index, query_param, index_dir,
|
||||
retrieval_mode, filter_mode);
|
||||
|
||||
cout << "Bench done." << endl;
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
index->Close();
|
||||
|
||||
return 0;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,147 @@
|
||||
// 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 <stdio.h>
|
||||
#include <string.h>
|
||||
#include <limits>
|
||||
#include <map>
|
||||
#ifdef _MSC_VER
|
||||
#include <chrono>
|
||||
#else
|
||||
#include <sys/time.h>
|
||||
#endif
|
||||
#include <ailego/parallel/lock.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class BenchResult {
|
||||
public:
|
||||
BenchResult() {
|
||||
total_query_count_ = 0;
|
||||
total_process_time_by_us_ = 0;
|
||||
min_time_by_us_ = std::numeric_limits<long>::max();
|
||||
max_time_by_us_ = 0;
|
||||
}
|
||||
~BenchResult() {}
|
||||
|
||||
void add_time(int query_count, long time_by_us) {
|
||||
lock_.lock();
|
||||
total_query_count_ += query_count;
|
||||
total_process_time_by_us_ += time_by_us;
|
||||
long time_val = time_by_us / 100;
|
||||
if (process_time_map_.find(time_val) != process_time_map_.end()) {
|
||||
++process_time_map_[time_val];
|
||||
} else {
|
||||
process_time_map_[time_val] = 1;
|
||||
}
|
||||
if (time_by_us < min_time_by_us_) {
|
||||
min_time_by_us_ = time_by_us;
|
||||
} else if (time_by_us > max_time_by_us_) {
|
||||
max_time_by_us_ = time_by_us;
|
||||
}
|
||||
lock_.unlock();
|
||||
}
|
||||
void mark_start() {
|
||||
#ifdef _MSC_VER
|
||||
start_ = std::chrono::steady_clock::now();
|
||||
#else
|
||||
gettimeofday(&start_, NULL);
|
||||
#endif
|
||||
}
|
||||
void mark_end() {
|
||||
#ifdef _MSC_VER
|
||||
end_ = std::chrono::steady_clock::now();
|
||||
#else
|
||||
gettimeofday(&end_, NULL);
|
||||
#endif
|
||||
}
|
||||
long get_duration_by_ms() {
|
||||
#ifdef _MSC_VER
|
||||
return static_cast<long>(
|
||||
std::chrono::duration_cast<std::chrono::milliseconds>(end_ - start_)
|
||||
.count());
|
||||
#else
|
||||
long duration = (end_.tv_sec - start_.tv_sec) * 1000 +
|
||||
(end_.tv_usec - start_.tv_usec) / 1000;
|
||||
return duration;
|
||||
#endif
|
||||
}
|
||||
long get_total_query_count() {
|
||||
return total_query_count_;
|
||||
}
|
||||
std::map<long, long> &get_process_time_map() {
|
||||
return process_time_map_;
|
||||
}
|
||||
long get_total_process_time_by_ms() {
|
||||
return total_process_time_by_us_ / 1000;
|
||||
}
|
||||
void print() {
|
||||
fprintf(stdout,
|
||||
"Process query: %ld, total process time: %ldms, "
|
||||
"duration: %ldms, max: %ldms, min:%ldms\n",
|
||||
get_total_query_count(), get_total_process_time_by_ms(),
|
||||
get_duration_by_ms(), max_time_by_us_ / 1000,
|
||||
min_time_by_us_ / 1000);
|
||||
fprintf(stdout, "Avg latency: %0.1fms qps: %0.1f\n",
|
||||
((float)get_total_process_time_by_ms()) / get_total_query_count(),
|
||||
get_total_query_count() / ((float)get_duration_by_ms() / 1000));
|
||||
|
||||
int tot_num = 0;
|
||||
int percent[] = {25, 50, 75, 90, 95, 99};
|
||||
int index = 0;
|
||||
float max_time = 0.0;
|
||||
int last_num = 0;
|
||||
|
||||
for (auto element : process_time_map_) {
|
||||
tot_num += element.second;
|
||||
if (tot_num >= total_query_count_ * percent[index] / 100) {
|
||||
if (last_num != tot_num) {
|
||||
max_time = (float)element.first / 10;
|
||||
last_num = tot_num;
|
||||
}
|
||||
fprintf(stdout, "%d Percentile:\t\t %.1f ms\n", percent[index],
|
||||
max_time);
|
||||
index++;
|
||||
if (index >= 6) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (; index < 6; index++) {
|
||||
fprintf(stdout, "%d Percentile:\t\t %.1f ms\n", percent[index], max_time);
|
||||
}
|
||||
fprintf(stdout, "\n");
|
||||
}
|
||||
|
||||
private:
|
||||
long total_query_count_;
|
||||
long total_process_time_by_us_;
|
||||
long min_time_by_us_;
|
||||
long max_time_by_us_;
|
||||
#ifdef _MSC_VER
|
||||
std::chrono::steady_clock::time_point start_;
|
||||
std::chrono::steady_clock::time_point end_;
|
||||
#else
|
||||
struct timeval start_;
|
||||
struct timeval end_;
|
||||
#endif
|
||||
ailego::SpinMutex lock_;
|
||||
std::map<long, long> process_time_map_; // <processTimeBy100us, count>
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
@@ -0,0 +1,187 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import polars as pl
|
||||
|
||||
to_append = True
|
||||
|
||||
|
||||
def write_neighbors_file(data_frame, neighbors_file):
|
||||
id_list = np.stack(data_frame["id"])
|
||||
neighbors_list = np.stack(data_frame["neighbors_id"])
|
||||
|
||||
id_list.tolist()
|
||||
neighbors_list.tolist()
|
||||
|
||||
if len(id_list) != len(neighbors_list):
|
||||
logger.error("list size not equal: %d, %d", len(id_list), len(neighbors_list))
|
||||
os._exit(1)
|
||||
|
||||
for i in range(len(id_list)):
|
||||
id_int = id_list[i]
|
||||
line = str(id_int) + ";"
|
||||
|
||||
neighbors = neighbors_list[i]
|
||||
# for j in range(len(neighbors)):
|
||||
for j in range(100):
|
||||
neighbor_id = neighbors[j]
|
||||
|
||||
line += str(neighbor_id)
|
||||
if j != 99:
|
||||
line += " "
|
||||
else:
|
||||
line += "\n"
|
||||
|
||||
neighbors_file.write(line)
|
||||
|
||||
logger.info("Output neighbors file done. Total lines: %d", len(id_list))
|
||||
|
||||
|
||||
def write_vector_file(data_frame, vector_file):
|
||||
test_embedding_list = np.stack(data_frame["emb"])
|
||||
test_embedding_list.tolist()
|
||||
|
||||
test_id_list = np.stack(data_frame["id"])
|
||||
test_id_list.tolist()
|
||||
|
||||
if len(test_id_list) != len(test_embedding_list):
|
||||
logger.info(
|
||||
"id list not matched with embedding list! : %d, %d",
|
||||
len(test_id_list),
|
||||
len(test_embedding_list),
|
||||
)
|
||||
return
|
||||
|
||||
for case_id in range(len(test_id_list)):
|
||||
idx = test_id_list[case_id]
|
||||
vector = test_embedding_list[case_id]
|
||||
|
||||
vector_line = str(idx) + ";"
|
||||
|
||||
for i in range(len(vector)):
|
||||
vector_line += str(round(vector[i], 16))
|
||||
if i != len(vector) - 1:
|
||||
vector_line += " "
|
||||
else:
|
||||
vector_line += ";"
|
||||
|
||||
vector_line += "\n"
|
||||
|
||||
vector_file.write(vector_line)
|
||||
|
||||
if case_id != 0 and case_id % 10000 == 0:
|
||||
logger.info("output lines: %d", case_id)
|
||||
|
||||
logger.info("Output vector file done. Total lines: %d", len(test_id_list))
|
||||
|
||||
|
||||
def read_parquet_file(file_name: str) -> pd.DataFrame:
|
||||
parquet_file = pathlib.Path(file_name)
|
||||
if not parquet_file.exists():
|
||||
logger.error("open error!")
|
||||
return pd.DataFrame()
|
||||
|
||||
try:
|
||||
return pl.read_parquet(parquet_file)
|
||||
except Exception:
|
||||
logger.error("open error! error file: %s", file_name)
|
||||
return pd.DataFrame()
|
||||
|
||||
|
||||
def gen_vector_files(input_dir, input_file_pattern, output_dir, output_file_name):
|
||||
input_file_list = list(Path(input_dir).rglob(input_file_pattern))
|
||||
|
||||
output_file_name_full = pathlib.Path(output_dir, output_file_name)
|
||||
|
||||
if not to_append and output_file_name_full.exists():
|
||||
logger.error("File exists! File name: %s", output_file_name_full)
|
||||
os._exit(1)
|
||||
|
||||
write_flag = "a" if to_append else "w"
|
||||
|
||||
with Path.open(output_file_name_full.resolve(), write_flag) as vector_file:
|
||||
for input_file in input_file_list:
|
||||
input_file_name = input_file.resolve()
|
||||
|
||||
logger.info(
|
||||
"Load the entire file into memory. File name: %s", input_file_name
|
||||
)
|
||||
data_set = read_parquet_file(input_file.resolve())
|
||||
logger.info("Read parquet file done. File name: %s", input_file_name)
|
||||
|
||||
if len(data_set) > 0:
|
||||
logger.info("Process parquet file. File name: %s", input_file_name)
|
||||
write_vector_file(data_set, vector_file)
|
||||
logger.info("Process parquet file done. File name: %s", input_file_name)
|
||||
|
||||
|
||||
def gen_neighbor_files(input_dir, input_file_pattern, output_dir, output_file_name):
|
||||
input_file_list = list(Path(input_dir).rglob(input_file_pattern))
|
||||
|
||||
output_file_name_full = pathlib.Path(output_dir, output_file_name)
|
||||
|
||||
if not to_append and output_file_name_full.exists():
|
||||
logger.error("File already exists. File name: %s", output_file_name_full)
|
||||
os._exit(1)
|
||||
|
||||
write_flag = "a" if to_append else "w"
|
||||
|
||||
with Path.open(output_file_name_full.resolve(), write_flag) as neighbor_file:
|
||||
for input_file in input_file_list:
|
||||
input_file_name = input_file.resolve()
|
||||
|
||||
logger.info(
|
||||
"Load the entire file into memory. File name: %s", input_file_name
|
||||
)
|
||||
data_set = read_parquet_file(input_file.resolve())
|
||||
logger.info("Read parquet file done. File name: %s", input_file_name)
|
||||
|
||||
if len(data_set) > 0:
|
||||
logger.info("Write parquet file. File name: %s", input_file_name)
|
||||
write_neighbors_file(data_set, neighbor_file)
|
||||
logger.info("Write parquet file done. File name: %s", input_file_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger = logging.getLogger("convert_log")
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.DEBUG)
|
||||
|
||||
formatter = logging.Formatter(
|
||||
fmt="%(asctime)s [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
|
||||
)
|
||||
|
||||
console_handler.setFormatter(formatter)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
input_dir = "./cohere/10m"
|
||||
output_dir = "./10m.output"
|
||||
|
||||
logger.info("Generate test vector files")
|
||||
input_file_pattern = "test.parquet"
|
||||
output_file_name = "cohere_test_vector_1000.new.txt"
|
||||
|
||||
to_append = False
|
||||
gen_vector_files(input_dir, input_file_pattern, output_dir, output_file_name)
|
||||
|
||||
logger.info("Generate neighbor files")
|
||||
input_file_pattern = "neighbors.parquet"
|
||||
output_file_name = "neighbors.txt"
|
||||
|
||||
to_append = False
|
||||
gen_neighbor_files(input_dir, input_file_pattern, output_dir, output_file_name)
|
||||
|
||||
logger.info("Generate train vector files")
|
||||
output_file_name = "cohere_768_10m_vector.train.txt"
|
||||
to_append = True
|
||||
for i in range(10):
|
||||
input_file_pattern = "shuffle_train-0" + str(i) + "-of-10.parquet"
|
||||
gen_vector_files(input_dir, input_file_pattern, output_dir, output_file_name)
|
||||
@@ -0,0 +1,77 @@
|
||||
// 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 <vector>
|
||||
#include <roaring/roaring.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
struct FilterResultCache {
|
||||
public:
|
||||
FilterResultCache() {
|
||||
bitmap_ = roaring_bitmap_create();
|
||||
}
|
||||
|
||||
FilterResultCache(uint32_t capacity_hint) {
|
||||
bitmap_ = roaring_bitmap_create_with_capacity(capacity_hint);
|
||||
}
|
||||
|
||||
~FilterResultCache() {
|
||||
roaring_bitmap_free(bitmap_);
|
||||
bitmap_ = nullptr;
|
||||
}
|
||||
|
||||
bool find(uint64_t key) const {
|
||||
return !roaring_bitmap_contains(bitmap_, key);
|
||||
}
|
||||
|
||||
void set(uint64_t key) const {
|
||||
roaring_bitmap_add(bitmap_, key);
|
||||
}
|
||||
|
||||
int filter(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
|
||||
const std::vector<uint64_t> &query_tag_list,
|
||||
const std::vector<uint64_t> &id_to_key_list) {
|
||||
for (size_t i = 0; i < id_to_tags_list.size(); ++i) {
|
||||
auto &id_tag_list = id_to_tags_list[i];
|
||||
|
||||
size_t t_i = 0;
|
||||
size_t q_i = 0;
|
||||
while (t_i < id_tag_list.size() && q_i < query_tag_list.size()) {
|
||||
if (id_tag_list[t_i] == query_tag_list[q_i]) {
|
||||
uint64_t key = id_to_key_list[i];
|
||||
|
||||
set(key);
|
||||
|
||||
break;
|
||||
} else if (id_tag_list[t_i] < query_tag_list[q_i]) {
|
||||
++t_i;
|
||||
} else {
|
||||
++q_i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
public:
|
||||
roaring_bitmap_t *bitmap_{nullptr};
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
@@ -0,0 +1,555 @@
|
||||
// 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
|
||||
@@ -0,0 +1,297 @@
|
||||
// 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 <sys/stat.h>
|
||||
#include <signal.h>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <mutex>
|
||||
#include <ailego/container/bitmap.h>
|
||||
#include <ailego/parallel/lock.h>
|
||||
#include <zvec/ailego/hash/crc32c.h>
|
||||
#include <zvec/ailego/io/file.h>
|
||||
#include <zvec/ailego/logger/logger.h>
|
||||
#include <zvec/ailego/parallel/thread_pool.h>
|
||||
#include <zvec/ailego/utility/string_helper.h>
|
||||
#include <zvec/ailego/utility/time_helper.h>
|
||||
#include "zvec/core/framework/index_error.h"
|
||||
#include "zvec/core/framework/index_factory.h"
|
||||
#include "zvec/core/framework/index_plugin.h"
|
||||
#include "zvec/core/framework/index_storage.h"
|
||||
#include "zvec/core/interface/index.h"
|
||||
#include "zvec/core/interface/index_factory.h"
|
||||
#include "zvec/core/interface/index_param.h"
|
||||
#include "filter_result_cache.h"
|
||||
#include "meta_segment_common.h"
|
||||
#include "txt_input_reader.h"
|
||||
|
||||
#ifdef __clang__
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wshadow"
|
||||
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
|
||||
#elif defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wshadow"
|
||||
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
|
||||
#endif
|
||||
|
||||
#include <yaml-cpp/yaml.h>
|
||||
|
||||
#ifdef __clang__
|
||||
#pragma clang diagnostic pop
|
||||
#elif defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
|
||||
using namespace std;
|
||||
using namespace zvec;
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
|
||||
int parse_and_load_index_param(
|
||||
YAML::Node &config_node, string &index_dir,
|
||||
core_interface::Index::Pointer &index,
|
||||
core_interface::BaseIndexQueryParam::Pointer &query_param) {
|
||||
// Create Index from config
|
||||
|
||||
if (auto index_config = config_node["IndexCommon"]["IndexConfig"]) {
|
||||
std::cout << "IndexConfig: " << index_config.as<string>() << std::endl;
|
||||
auto params = core_interface::IndexFactory::DeserializeIndexParamFromJson(
|
||||
index_config.as<string>());
|
||||
index = core_interface::IndexFactory::CreateAndInitIndex(*params);
|
||||
if (!index) {
|
||||
LOG_ERROR("Failed to create index");
|
||||
return -1;
|
||||
}
|
||||
core_interface::StorageOptions storage_options;
|
||||
storage_options.type = core_interface::StorageOptions::StorageType::kMMAP;
|
||||
storage_options.create_new = false;
|
||||
storage_options.read_only = true;
|
||||
|
||||
int ret = index->Open(index_dir, storage_options);
|
||||
if (0 != ret) {
|
||||
LOG_ERROR("Index open failed with ret %d", ret);
|
||||
return -1;
|
||||
}
|
||||
|
||||
cout << "Load index done!" << endl;
|
||||
} else {
|
||||
LOG_ERROR("IndexCommon.IndexConfig is required");
|
||||
return -1;
|
||||
}
|
||||
|
||||
/*
|
||||
QueryConfig:
|
||||
QueryParam: '{"ef_search":100,"index_type":"kHNSW"}'
|
||||
RefinerConfig:
|
||||
ScaleFactor: !!int 2
|
||||
ReferenceIndex:
|
||||
Config:
|
||||
'{"use_id_map":false,"data_type":"DT_FP32","dimension":768,"index_type":"kHNSW","metric_type":"kCosine"}'
|
||||
Path: ./cohere_train_vector_1m.2.index
|
||||
*/
|
||||
|
||||
// QUERY PARAM
|
||||
if (auto query_config = config_node["QueryConfig"]; query_config) {
|
||||
// QueryConfig.QueryParam
|
||||
if (auto query_param_config = query_config["QueryParam"];
|
||||
query_param_config) {
|
||||
std::cout << "QueryParam: " << query_param_config.as<string>()
|
||||
<< std::endl;
|
||||
query_param = core_interface::IndexFactory::QueryParamDeserializeFromJson<
|
||||
core_interface::BaseIndexQueryParam>(
|
||||
query_param_config.as<std::string>());
|
||||
if (!query_param) {
|
||||
LOG_ERROR("Failed to deserialize query params");
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
// QueryConfig.RefinerConfig
|
||||
if (auto refiner_config = query_config["RefinerConfig"]; refiner_config) {
|
||||
core_interface::Index::Pointer reference_index = nullptr;
|
||||
auto refiner_param = std::make_shared<core_interface::RefinerParam>();
|
||||
|
||||
// QueryConfig.RefinerConfig.ScaleFactor
|
||||
if (auto scale_factor_config = refiner_config["ScaleFactor"];
|
||||
scale_factor_config) {
|
||||
auto scale_factor = scale_factor_config.as<float>();
|
||||
refiner_param->scale_factor_ = scale_factor;
|
||||
} else {
|
||||
LOG_ERROR("QueryConfig.RefinerConfig.ScaleFactor config is required");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// QueryConfig.RefinerConfig.ReferenceIndex
|
||||
if (auto reference_index_config = refiner_config["ReferenceIndex"];
|
||||
reference_index_config) {
|
||||
// QueryConfig.RefinerConfig.ReferenceIndex.Config
|
||||
if (auto reference_index_config_config =
|
||||
reference_index_config["Config"];
|
||||
reference_index_config_config) {
|
||||
auto params =
|
||||
core_interface::IndexFactory::DeserializeIndexParamFromJson(
|
||||
reference_index_config_config.as<std::string>());
|
||||
|
||||
reference_index =
|
||||
core_interface::IndexFactory::CreateAndInitIndex(*params);
|
||||
} else {
|
||||
LOG_ERROR(
|
||||
"QueryConfig.RefinerConfig.ReferenceIndex.Config config is "
|
||||
"required");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// QueryConfig.RefinerConfig.ReferenceIndex.Path
|
||||
if (auto reference_index_path_config = reference_index_config["Path"];
|
||||
reference_index_path_config) {
|
||||
auto reference_index_path =
|
||||
reference_index_path_config.as<std::string>();
|
||||
core_interface::StorageOptions storage_options;
|
||||
storage_options.type =
|
||||
core_interface::StorageOptions::StorageType::kMMAP;
|
||||
storage_options.create_new = false;
|
||||
storage_options.read_only = true;
|
||||
|
||||
int ret =
|
||||
reference_index->Open(reference_index_path, storage_options);
|
||||
if (0 != ret) {
|
||||
LOG_ERROR("Index open failed with ret %d", ret);
|
||||
return -1;
|
||||
}
|
||||
|
||||
cout << "Load reference index done!" << endl;
|
||||
} else {
|
||||
LOG_ERROR(
|
||||
"QueryConfig.RefinerConfig.ReferenceIndex.Path is required");
|
||||
return -1;
|
||||
}
|
||||
refiner_param->reference_index = reference_index;
|
||||
} else {
|
||||
LOG_ERROR(
|
||||
"QueryConfig.RefinerConfig.ReferenceIndex section is required");
|
||||
return -1;
|
||||
} // QueryConfig.RefinerConfig.ReferenceIndex
|
||||
|
||||
query_param->refiner_param = refiner_param;
|
||||
} // QueryConfig.RefinerConfig
|
||||
} // QUERY PARAM
|
||||
return 0;
|
||||
}
|
||||
|
||||
//--------------------------------------------------
|
||||
// Helper functions for loading tag lists
|
||||
//--------------------------------------------------
|
||||
int load_taglists(const std::string &path,
|
||||
std::vector<std::vector<uint64_t>> &id_to_tags_list,
|
||||
std::vector<uint64_t> &tag_key_list) {
|
||||
// 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;
|
||||
tag_list.reserve(tag_count);
|
||||
for (size_t j = 0; j < tag_count; ++j) {
|
||||
tag_list.push_back(reinterpret_cast<const uint64_t *>(data_ptr)[j]);
|
||||
}
|
||||
id_to_tags_list.push_back(std::move(tag_list));
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
// 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 <string>
|
||||
#include "zvec/core/framework/index_meta.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class IndexMetaHelper {
|
||||
public:
|
||||
static std::string to_string(IndexMeta::DataType type) {
|
||||
switch (type) {
|
||||
case IndexMeta::DataType::DT_FP32:
|
||||
return std::string("FP32");
|
||||
case IndexMeta::DataType::DT_FP64:
|
||||
return std::string("FP64");
|
||||
case IndexMeta::DataType::DT_INT16:
|
||||
return std::string("INT16");
|
||||
case IndexMeta::DataType::DT_INT8:
|
||||
return std::string("INT8");
|
||||
case IndexMeta::DataType::DT_BINARY32:
|
||||
return std::string("Binary");
|
||||
case IndexMeta::DataType::DT_BINARY64:
|
||||
return std::string("Binary64");
|
||||
case IndexMeta::DataType::DT_FP16:
|
||||
return std::string("FP16");
|
||||
default:
|
||||
return std::string("NotSupportedType");
|
||||
}
|
||||
}
|
||||
|
||||
static std::string to_string(IndexMeta meta) {
|
||||
char buffer[1024];
|
||||
snprintf(buffer, 1024,
|
||||
"IndexMeta: type[%s] method[%s] dimension[%u] element_size[%u]",
|
||||
to_string(meta.data_type()).c_str(), meta.metric_name().c_str(),
|
||||
meta.dimension(), meta.element_size());
|
||||
return std::string(buffer);
|
||||
}
|
||||
|
||||
static bool parse_from(const std::string &type, const std::string &method,
|
||||
const std::string &vector_type, IndexMeta &meta) {
|
||||
return parse_from(type, method, 0, vector_type, meta);
|
||||
}
|
||||
|
||||
static bool parse_from(const std::string &type, const std::string &method,
|
||||
const size_t dimension, const std::string &vector_type,
|
||||
IndexMeta &meta) {
|
||||
if (vector_type != "dense" && vector_type != "sparse") {
|
||||
std::cerr << "vector type should be dense or sparse!!!" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
auto feature_type = IndexMeta::DataType::DT_UNDEFINED;
|
||||
if (type == std::string("float")) {
|
||||
feature_type = IndexMeta::DataType::DT_FP32;
|
||||
} else if (type == std::string("double")) {
|
||||
feature_type = IndexMeta::DataType::DT_FP64;
|
||||
} else if (type == std::string("int16")) {
|
||||
feature_type = IndexMeta::DataType::DT_INT16;
|
||||
} else if (type == std::string("int8")) {
|
||||
feature_type = IndexMeta::DataType::DT_INT8;
|
||||
} else if (type == std::string("binary")) {
|
||||
feature_type = IndexMeta::DataType::DT_BINARY32;
|
||||
} else if (type == std::string("binary64")) {
|
||||
feature_type = IndexMeta::DataType::DT_BINARY64;
|
||||
} else {
|
||||
std::cerr << "Not supported type: " << type << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
meta.set_meta(feature_type, dimension);
|
||||
ailego::Params params;
|
||||
if (method == std::string("L2")) {
|
||||
if (feature_type == IndexMeta::DataType::DT_FP32) {
|
||||
meta.set_metric("SquaredEuclidean", 0, std::move(params));
|
||||
} else if (feature_type == IndexMeta::DataType::DT_INT8) {
|
||||
meta.set_metric("SquaredEuclidean", 0, std::move(params));
|
||||
} else if (feature_type == IndexMeta::DataType::DT_FP16) {
|
||||
meta.set_metric("SquaredEuclidean", 0, std::move(params));
|
||||
} else {
|
||||
std::cerr << "Not supported type(" << type << ") for L2" << std::endl;
|
||||
return false;
|
||||
}
|
||||
} else if (method == std::string("IP")) {
|
||||
if (feature_type == IndexMeta::DataType::DT_FP32) {
|
||||
meta.set_metric("InnerProduct", 0, std::move(params));
|
||||
} else if (feature_type == IndexMeta::DataType::DT_INT8) {
|
||||
meta.set_metric("InnerProduct", 0, std::move(params));
|
||||
} else if (feature_type == IndexMeta::DataType::DT_FP16) {
|
||||
meta.set_metric("InnerProduct", 0, std::move(params));
|
||||
} else {
|
||||
std::cerr << "Not supported type(" << type << ") for IP" << std::endl;
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
std::cerr << "Not supported method: " << method << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,34 @@
|
||||
// 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/ailego/utility/type_helper.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
const static std::string TAGLIST_HEADER_SEGMENT_NAME("local_taglists_header");
|
||||
const static std::string TAGLIST_KEY_SEGMENT_NAME("local_taglists_key");
|
||||
const static std::string TAGLIST_DATA_SEGMENT_NAME("local_taglists_data");
|
||||
|
||||
#pragma pack(4)
|
||||
struct TagListHeader {
|
||||
uint64_t num_vecs;
|
||||
uint8_t meta_buf[252];
|
||||
};
|
||||
#pragma pack()
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,723 @@
|
||||
// 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 <iostream>
|
||||
#include <set>
|
||||
#include "gflags/gflags.h"
|
||||
#include "zvec/core/framework/index_meta.h"
|
||||
#include "index_meta_helper.h"
|
||||
#include "txt_input_reader.h"
|
||||
#include "vecs_common.h"
|
||||
|
||||
using namespace std;
|
||||
using namespace zvec::core;
|
||||
|
||||
DEFINE_string(input, "input.txt", "txt input file");
|
||||
DEFINE_string(input_first_sep, ";", "input first sep");
|
||||
DEFINE_string(input_second_sep, " ", "input second sep");
|
||||
DEFINE_string(output, "output.vecs", "vecs output file");
|
||||
DEFINE_string(type, "float",
|
||||
"available type: float, double, int16, int8, binary");
|
||||
DEFINE_string(method, "L2", "available method: L2, IP");
|
||||
DEFINE_int32(dimension, 256, "data dimension");
|
||||
DEFINE_string(vector_type, "dense", "available type: dense, hybrid, sparse");
|
||||
|
||||
bool write_header_output(VecsHeader &header, const IndexMeta &meta,
|
||||
size_t &total_writes, FILE *wfp) {
|
||||
// write header
|
||||
std::cout << "Begin to Write Header Section..." << std::endl;
|
||||
|
||||
std::string meta_buf;
|
||||
meta.serialize(&meta_buf);
|
||||
header.meta_size = meta_buf.size();
|
||||
size_t wret = fwrite(&header, sizeof(header), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write header error" << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(header);
|
||||
std::cout << "Total Writes after Header Section: " << total_writes
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// write meta
|
||||
std::cout << "Begin to Write Meta Section..." << std::endl;
|
||||
wret = fwrite(meta_buf.c_str(), meta_buf.size(), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write header meta_buf error" << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += meta_buf.size();
|
||||
std::cout << "Total Writes after Meta Buf: " << total_writes << std::endl
|
||||
<< std::endl;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool write_header_output_sparse(VecsHeader &header, const IndexMeta &meta,
|
||||
size_t &total_writes, FILE *wfp) {
|
||||
// write header
|
||||
std::cout << "Begin to Write Header Section..." << std::endl;
|
||||
std::string meta_buf;
|
||||
meta.serialize(&meta_buf);
|
||||
header.meta_size = meta_buf.size();
|
||||
size_t wret = fwrite(&header, sizeof(header), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write header error" << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(header);
|
||||
std::cout << "Total Writes after Header Section: " << total_writes
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// write meta
|
||||
std::cout << "Begin to Write Sparse Meta Section..." << std::endl;
|
||||
wret = fwrite(meta_buf.c_str(), meta_buf.size(), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write header meta buf error" << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += meta_buf.size();
|
||||
std::cout << "Total Writes after Meta Buf: " << total_writes << std::endl
|
||||
<< std::endl;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool write_features_output(size_t vec_num, const vector<vector<T>> &features,
|
||||
size_t &total_writes, FILE *wfp) {
|
||||
// write dense vector
|
||||
std::cout << "Begin to Write Dense Vector Section..." << std::endl;
|
||||
for (size_t i = 0; i < vec_num; ++i) {
|
||||
auto &feature = features[i];
|
||||
size_t wret = fwrite(&feature[0], sizeof(T), feature.size(), wfp);
|
||||
if (wret != feature.size()) {
|
||||
cerr << "Write feature error. " << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += feature.size() * sizeof(T);
|
||||
}
|
||||
|
||||
std::cout << "Total Writes after Dense Vector: " << total_writes << std::endl
|
||||
<< std::endl;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool write_keys_output(size_t vec_num, const vector<uint64_t> &keys,
|
||||
size_t &total_writes, FILE *wfp) {
|
||||
std::cout << "Begin to Write Key Section..." << std::endl;
|
||||
for (size_t i = 0; i < vec_num; ++i) {
|
||||
uint64_t key = keys[i];
|
||||
size_t wret = fwrite(&key, sizeof(key), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write key error. key:" << key << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(uint64_t);
|
||||
}
|
||||
|
||||
std::cout << "Total Writes after Key Section: " << total_writes << std::endl
|
||||
<< std::endl;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool write_sparse_features_output(size_t vec_num,
|
||||
const vector<SparseData<T>> &sparse_data,
|
||||
size_t &total_writes, FILE *wfp) {
|
||||
std::set<uint32_t> sparse_dims;
|
||||
uint32_t sparse_max_count = 0;
|
||||
uint32_t sparse_min_count = -1U;
|
||||
uint32_t sparse_total_count = 0;
|
||||
|
||||
// write sparse meta
|
||||
std::cout << "Begin to Write Sparse Meta Section..." << std::endl;
|
||||
size_t wret;
|
||||
uint64_t offset = 0;
|
||||
for (size_t i = 0; i < vec_num; ++i) {
|
||||
wret = fwrite(&offset, sizeof(uint64_t), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write sparse feature len error. " << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
offset += sparse_data[i].get_len();
|
||||
|
||||
total_writes += sizeof(size_t);
|
||||
}
|
||||
std::cout << "Total Writes after Sparse Meta Section: " << total_writes
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
std::cout << "Begin to Write Sparse Vector Section..." << std::endl;
|
||||
for (size_t i = 0; i < vec_num; ++i) {
|
||||
auto &sparse_one_data = sparse_data[i];
|
||||
|
||||
wret = fwrite(&(sparse_one_data.count), sizeof(uint32_t), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write sparse feature count error. " << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(uint32_t);
|
||||
|
||||
wret = fwrite(&sparse_one_data.indices[0], sizeof(uint32_t),
|
||||
sparse_one_data.indices.size(), wfp);
|
||||
if (wret != sparse_one_data.indices.size()) {
|
||||
cerr << "Write feature error. " << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(uint32_t) * sparse_one_data.indices.size();
|
||||
// do some stat
|
||||
for (size_t s = 0; s < sparse_one_data.indices.size(); ++s) {
|
||||
sparse_dims.insert(sparse_one_data.indices[s]);
|
||||
}
|
||||
|
||||
if (sparse_one_data.indices.size() > sparse_max_count) {
|
||||
sparse_max_count = sparse_one_data.indices.size();
|
||||
}
|
||||
|
||||
if (sparse_one_data.indices.size() < sparse_min_count) {
|
||||
sparse_min_count = sparse_one_data.indices.size();
|
||||
}
|
||||
|
||||
sparse_total_count += sparse_one_data.indices.size();
|
||||
// //done
|
||||
|
||||
wret = fwrite(&sparse_one_data.features[0], sizeof(T),
|
||||
sparse_one_data.features.size(), wfp);
|
||||
if (wret != sparse_one_data.features.size()) {
|
||||
cerr << "Write feature error. " << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(T) * sparse_one_data.features.size();
|
||||
}
|
||||
|
||||
std::cout << "Total Writes after Sparse Vector Section: " << total_writes
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
// for (auto itr=sparse_dims.begin(); itr!=sparse_dims.end(); ++itr) {
|
||||
// std::cout << (*itr) << ",";
|
||||
// }
|
||||
// std::cout << std::endl;
|
||||
|
||||
std::cout << "Max Sparse Dimension Count: " << sparse_max_count << std::endl;
|
||||
std::cout << "Min Sparse Dimension Count: " << sparse_min_count << std::endl;
|
||||
std::cout << "Avg Sparse Dimension Count: " << sparse_total_count / vec_num
|
||||
<< std::endl;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool write_taglists_output(size_t vec_num,
|
||||
const vector<vector<uint64_t>> &taglists,
|
||||
size_t &total_writes, FILE *wfp) {
|
||||
std::cout << "Begin to Write Tag List Section..." << std::endl;
|
||||
|
||||
// write tag list meta
|
||||
std::cout << "Begin to Write Tag List Meta Section..." << std::endl;
|
||||
size_t wret;
|
||||
uint64_t offset = 0;
|
||||
for (size_t i = 0; i < vec_num; ++i) {
|
||||
wret = fwrite(&offset, sizeof(uint64_t), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write tag list meta error. Rec no: " << i << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
offset += taglists[i].size() * sizeof(uint64_t);
|
||||
|
||||
total_writes += sizeof(size_t);
|
||||
}
|
||||
std::cout << "Total Writes after Tag Meta Section: " << total_writes
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
for (size_t i = 0; i < vec_num; ++i) {
|
||||
std::vector<uint64_t> taglist = taglists[i];
|
||||
uint64_t taglist_size = taglist.size();
|
||||
wret = fwrite(&taglist_size, sizeof(uint64_t), 1, wfp);
|
||||
if (wret != 1) {
|
||||
cerr << "Write tag list size error. Rec no: " << i << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
wret = fwrite(&(taglist[0]), sizeof(uint64_t), taglist.size(), wfp);
|
||||
if (wret != taglist.size()) {
|
||||
cerr << "Write tag list error. Rec no: " << i << endl;
|
||||
fclose(wfp);
|
||||
return false;
|
||||
}
|
||||
|
||||
total_writes += sizeof(uint64_t) * taglist.size() + sizeof(uint64_t);
|
||||
}
|
||||
|
||||
std::cout << "Total Writes after Tag List Section: " << total_writes
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool write_vecs_output_sparse(VecsHeader &header, const IndexMeta &meta,
|
||||
const vector<uint64_t> &keys,
|
||||
const vector<SparseData<T>> &sparse_data,
|
||||
const vector<vector<uint64_t>> &taglists) {
|
||||
if (keys.empty()) {
|
||||
cerr << "keys is empty." << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
if (keys.size() != sparse_data.size()) {
|
||||
cerr << "keys's size(" << keys.size()
|
||||
<< ") is not equal to sparse data's size(" << sparse_data.size()
|
||||
<< ")." << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t vec_num = keys.size();
|
||||
|
||||
FILE *wfp = fopen(FLAGS_output.c_str(), "wb");
|
||||
if (!wfp) {
|
||||
cerr << "Open file error. " << FLAGS_output << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t total_writes = 0;
|
||||
|
||||
std::cout << "------------------------" << std::endl;
|
||||
std::cout << " Output Process " << std::endl;
|
||||
std::cout << "------------------------" << std::endl;
|
||||
|
||||
// write sparse header
|
||||
bool ret = write_header_output_sparse(header, meta, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write header error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
// write keys
|
||||
ret = write_keys_output(vec_num, keys, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write keys error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
// write sparse features
|
||||
ret = write_sparse_features_output(vec_num, sparse_data, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write sparse features error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
if ((header.bitmap & (1ULL << BITMAP_INDEX_TAGLIST)) != 0) {
|
||||
// write tag lists features
|
||||
ret = write_taglists_output(vec_num, taglists, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write tag lists error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
std::cout << "------------------------" << std::endl;
|
||||
std::cout << " Output Done " << std::endl;
|
||||
std::cout << "------------------------" << std::endl;
|
||||
|
||||
fclose(wfp);
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool write_vecs_output(VecsHeader &header, const IndexMeta &meta,
|
||||
const vector<uint64_t> &keys,
|
||||
const vector<vector<T>> &features,
|
||||
const vector<SparseData<T>> &sparse_data,
|
||||
const vector<vector<uint64_t>> &taglists) {
|
||||
if (keys.empty()) {
|
||||
cerr << "keys is empty." << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
if (keys.size() != features.size()) {
|
||||
cerr << "keys's size(" << keys.size()
|
||||
<< ") is not equal to features's size(" << features.size() << ")."
|
||||
<< endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
|
||||
size_t vec_num = header.num_vecs;
|
||||
|
||||
FILE *wfp = fopen(FLAGS_output.c_str(), "wb");
|
||||
if (!wfp) {
|
||||
cerr << "Open file error. " << FLAGS_output << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t total_writes = 0;
|
||||
|
||||
std::cout << "------------------------" << std::endl;
|
||||
std::cout << " Output Process " << std::endl;
|
||||
std::cout << "------------------------" << std::endl;
|
||||
|
||||
// write header
|
||||
bool ret = write_header_output(header, meta, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write header error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
// write features
|
||||
ret = write_features_output(vec_num, features, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write features error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
// write keys
|
||||
ret = write_keys_output(vec_num, keys, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write keys error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
// write sparse features
|
||||
if ((header.bitmap & (1ULL << BITMAP_INDEX_SPARSE)) != 0) {
|
||||
ret = write_sparse_features_output(vec_num, sparse_data, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write sparse features error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if ((header.bitmap & (1ULL << BITMAP_INDEX_TAGLIST)) != 0) {
|
||||
// write tag lists features
|
||||
ret = write_taglists_output(vec_num, taglists, total_writes, wfp);
|
||||
if (!ret) {
|
||||
cerr << "write tag lists error! " << endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
std::cout << "------------------------" << std::endl;
|
||||
std::cout << " Output Done " << std::endl;
|
||||
std::cout << "------------------------" << std::endl;
|
||||
|
||||
fclose(wfp);
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool compute_offset(uint64_t num_vecs, const IndexMeta &meta,
|
||||
const vector<uint64_t> & /*keys*/,
|
||||
const vector<vector<T>> & /*features*/,
|
||||
const vector<SparseData<T>> &sparse_data,
|
||||
const vector<std::vector<uint64_t>> &taglists,
|
||||
uint64_t &key_offset, uint64_t &feature_offset,
|
||||
uint64_t &sparse_offset, uint64_t &taglist_offset,
|
||||
uint64_t &key_size, uint64_t &feature_size,
|
||||
uint64_t &sparse_size, uint64_t &taglist_size) {
|
||||
size_t total_offset = 0;
|
||||
|
||||
feature_offset = 0;
|
||||
feature_size = num_vecs * meta.element_size();
|
||||
total_offset += feature_size;
|
||||
|
||||
key_offset = total_offset;
|
||||
key_size = num_vecs * sizeof(uint64_t);
|
||||
total_offset += key_size;
|
||||
|
||||
if (sparse_data.size() != 0) {
|
||||
sparse_offset = total_offset;
|
||||
|
||||
size_t data_offset = num_vecs * sizeof(uint64_t);
|
||||
for (size_t i = 0; i < sparse_data.size(); ++i) {
|
||||
data_offset += sizeof(uint32_t) +
|
||||
sparse_data[i].count * (sizeof(uint32_t) + sizeof(T));
|
||||
}
|
||||
|
||||
sparse_size = data_offset;
|
||||
|
||||
total_offset += sparse_size;
|
||||
} else {
|
||||
sparse_offset = -1LLU;
|
||||
sparse_size = 0;
|
||||
}
|
||||
|
||||
if (taglists.size() != 0) {
|
||||
taglist_offset = total_offset;
|
||||
|
||||
size_t data_offset = num_vecs * sizeof(uint64_t);
|
||||
for (size_t i = 0; i < taglists.size(); ++i) {
|
||||
data_offset += sizeof(uint64_t) + taglists[i].size() * sizeof(uint64_t);
|
||||
}
|
||||
|
||||
taglist_size = data_offset;
|
||||
} else {
|
||||
taglist_offset = -1LLU;
|
||||
taglist_size = 0;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool compute_sparse_offset(uint64_t num_vecs, const IndexMeta & /*meta*/,
|
||||
const vector<uint64_t> & /*keys*/,
|
||||
const vector<SparseData<T>> &sparse_data,
|
||||
const vector<std::vector<uint64_t>> &taglists,
|
||||
uint64_t &key_offset, uint64_t &sparse_offset,
|
||||
uint64_t &taglist_offset, uint64_t &key_size,
|
||||
uint64_t &sparse_size, uint64_t &taglist_size) {
|
||||
size_t total_offset = 0;
|
||||
|
||||
key_offset = 0;
|
||||
key_size = num_vecs * sizeof(uint64_t);
|
||||
total_offset += num_vecs * sizeof(uint64_t);
|
||||
|
||||
sparse_offset = total_offset;
|
||||
size_t data_offset = num_vecs * sizeof(uint64_t);
|
||||
for (size_t i = 0; i < sparse_data.size(); ++i) {
|
||||
data_offset += sizeof(uint32_t) +
|
||||
sparse_data[i].count * (sizeof(uint32_t) + sizeof(T));
|
||||
}
|
||||
|
||||
sparse_size = data_offset;
|
||||
total_offset += sparse_size;
|
||||
|
||||
if (taglists.size() != 0) {
|
||||
taglist_offset = total_offset;
|
||||
|
||||
data_offset = num_vecs * sizeof(uint64_t);
|
||||
for (size_t i = 0; i < taglists.size(); ++i) {
|
||||
data_offset += sizeof(uint64_t) + taglists[i].size() * sizeof(uint64_t);
|
||||
}
|
||||
|
||||
taglist_size = data_offset;
|
||||
} else {
|
||||
taglist_offset = -1LLU;
|
||||
taglist_size = 0;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool process(void) {
|
||||
if (FLAGS_vector_type == "sparse") {
|
||||
std::cout << "------------------------" << std::endl;
|
||||
std::cout << " Vector Type: sparse " << std::endl;
|
||||
std::cout << "------------------------" << std::endl;
|
||||
|
||||
IndexMeta meta;
|
||||
if (!IndexMetaHelper::parse_from(FLAGS_type, FLAGS_method,
|
||||
FLAGS_vector_type, meta)) {
|
||||
cerr << "Index meta parse error." << endl;
|
||||
return false;
|
||||
}
|
||||
cerr << IndexMetaHelper::to_string(meta) << endl;
|
||||
|
||||
TxtInputReader<T> reader;
|
||||
vector<uint64_t> keys;
|
||||
vector<SparseData<T>> sparse_data;
|
||||
vector<std::vector<uint64_t>> taglists;
|
||||
|
||||
bool ret = reader.load_record_sparse(FLAGS_input, FLAGS_input_first_sep,
|
||||
FLAGS_input_second_sep, keys,
|
||||
sparse_data, taglists);
|
||||
if (!ret) {
|
||||
cerr << "Read record failed" << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
if (sparse_data.size() == 0) {
|
||||
cerr << "empty sparse data!" << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t num_vecs = keys.size();
|
||||
|
||||
uint64_t key_offset{-1LLU}, sparse_offset{-1LLU}, taglist_offset{-1LLU};
|
||||
uint64_t key_size{0}, sparse_size{0}, taglist_size{0};
|
||||
|
||||
compute_sparse_offset(num_vecs, meta, keys, sparse_data, taglists,
|
||||
key_offset, sparse_offset, taglist_offset, key_size,
|
||||
sparse_size, taglist_size);
|
||||
|
||||
VecsHeader header;
|
||||
header.num_vecs = keys.size();
|
||||
header.meta_size_v1 = 0;
|
||||
header.version = 1;
|
||||
header.bitmap = 0;
|
||||
header.key_offset = key_offset;
|
||||
header.dense_offset = -1LLU;
|
||||
header.sparse_offset = sparse_offset;
|
||||
header.taglist_offset = taglist_offset;
|
||||
header.key_size = key_size;
|
||||
header.dense_size = 0;
|
||||
header.sparse_size = sparse_size;
|
||||
header.taglist_size = taglist_size;
|
||||
|
||||
header.bitmap |= (1 << BITMAP_INDEX_KEY);
|
||||
header.bitmap |= (1 << BITMAP_INDEX_SPARSE);
|
||||
|
||||
if (taglist_offset != -1LLU) {
|
||||
header.bitmap |= (1 << BITMAP_INDEX_TAGLIST);
|
||||
}
|
||||
|
||||
ret = write_vecs_output_sparse(header, meta, keys, sparse_data, taglists);
|
||||
if (!ret) {
|
||||
cerr << "write vecs output failed" << endl;
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
std::cout << "------------------------" << std::endl;
|
||||
std::cout << " Vector Type: " << FLAGS_vector_type << std::endl;
|
||||
std::cout << "------------------------" << std::endl;
|
||||
|
||||
IndexMeta meta;
|
||||
if (!IndexMetaHelper::parse_from(FLAGS_type, FLAGS_method, FLAGS_dimension,
|
||||
FLAGS_vector_type, meta)) {
|
||||
cerr << "Index meta parse error." << endl;
|
||||
return false;
|
||||
}
|
||||
cerr << IndexMetaHelper::to_string(meta) << endl;
|
||||
|
||||
TxtInputReader<T> reader;
|
||||
vector<uint64_t> keys;
|
||||
vector<vector<T>> features;
|
||||
vector<SparseData<T>> sparse_data;
|
||||
vector<std::vector<uint64_t>> taglists;
|
||||
|
||||
bool ret = reader.load_record(FLAGS_input, FLAGS_input_first_sep,
|
||||
FLAGS_input_second_sep, FLAGS_dimension, keys,
|
||||
features, sparse_data, taglists);
|
||||
if (!ret) {
|
||||
cerr << "Read record failed" << endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t num_vecs = keys.size();
|
||||
|
||||
uint64_t key_offset{-1LLU}, features_offset{-1LLU}, sparse_offset{-1LLU},
|
||||
taglist_offset{-1LLU};
|
||||
uint64_t key_size{0}, feature_size{0}, sparse_size{0}, taglist_size{0};
|
||||
|
||||
compute_offset(num_vecs, meta, keys, features, sparse_data, taglists,
|
||||
key_offset, features_offset, sparse_offset, taglist_offset,
|
||||
key_size, feature_size, sparse_size, taglist_size);
|
||||
|
||||
VecsHeader header;
|
||||
header.num_vecs = num_vecs;
|
||||
header.meta_size_v1 = 0;
|
||||
header.version = 1;
|
||||
header.bitmap = 0;
|
||||
header.key_offset = key_offset;
|
||||
header.dense_offset = features_offset;
|
||||
header.sparse_offset = sparse_offset;
|
||||
header.taglist_offset = taglist_offset;
|
||||
header.key_size = key_size;
|
||||
header.dense_size = feature_size;
|
||||
header.sparse_size = sparse_size;
|
||||
header.taglist_size = taglist_size;
|
||||
|
||||
header.bitmap |= (1 << BITMAP_INDEX_KEY);
|
||||
header.bitmap |= (1 << BITMAP_INDEX_DENSE);
|
||||
|
||||
if (sparse_offset != -1LLU) {
|
||||
header.bitmap |= (1 << BITMAP_INDEX_SPARSE);
|
||||
}
|
||||
|
||||
if (taglist_offset != -1LLU) {
|
||||
header.bitmap |= (1 << BITMAP_INDEX_TAGLIST);
|
||||
}
|
||||
|
||||
ret =
|
||||
write_vecs_output(header, meta, keys, features, sparse_data, taglists);
|
||||
if (!ret) {
|
||||
cerr << "write vecs output failed" << endl;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
int main(int argc, char *argv[]) {
|
||||
// gflags
|
||||
gflags::SetUsageMessage("Usage: txt2vecs [options]");
|
||||
gflags::ParseCommandLineFlags(&argc, &argv, true);
|
||||
|
||||
if (FLAGS_type == "float") {
|
||||
if (!process<float>()) {
|
||||
return -1;
|
||||
}
|
||||
} else if (FLAGS_type == "double") {
|
||||
if (!process<double>()) {
|
||||
return -1;
|
||||
}
|
||||
} else if (FLAGS_type == "int16") {
|
||||
if (!process<int16_t>()) {
|
||||
return -1;
|
||||
}
|
||||
} else if (FLAGS_type == "int8") {
|
||||
if (!process<int8_t>()) {
|
||||
return -1;
|
||||
}
|
||||
} else if (FLAGS_type == "binary") {
|
||||
if (!process<uint32_t>()) {
|
||||
return -1;
|
||||
}
|
||||
} else if (FLAGS_type == "binary64") {
|
||||
if (!process<uint64_t>()) {
|
||||
return -1;
|
||||
}
|
||||
} else {
|
||||
cerr << "Can not recognize type: " << FLAGS_type << endl;
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,549 @@
|
||||
// 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 <string.h>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <zvec/ailego/utility/string_helper.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
template <typename T>
|
||||
struct SparseData {
|
||||
public:
|
||||
SparseData(uint32_t count_in, std::vector<uint32_t> &indices_in,
|
||||
std::vector<T> &features_in)
|
||||
: count(count_in),
|
||||
indices(std::move(indices_in)),
|
||||
features(std::move(features_in)) {}
|
||||
|
||||
SparseData(uint32_t count_in, std::vector<uint32_t> &&indices_in,
|
||||
std::vector<T> &&features_in)
|
||||
: count(count_in),
|
||||
indices(std::move(indices_in)),
|
||||
features(std::move(features_in)) {}
|
||||
|
||||
public:
|
||||
uint64_t get_len() const {
|
||||
return sizeof(uint32_t) + sizeof(uint32_t) * indices.size() +
|
||||
sizeof(T) * features.size();
|
||||
}
|
||||
|
||||
public:
|
||||
uint32_t count;
|
||||
std::vector<uint32_t> indices;
|
||||
std::vector<T> features;
|
||||
};
|
||||
|
||||
// support type: float, binary, int16, int8
|
||||
template <typename T>
|
||||
class TxtInputReader {
|
||||
public:
|
||||
bool load_query(const std::string &query_file, const std::string &first_sep,
|
||||
const std::string &second_sep,
|
||||
std::vector<std::vector<T>> &features,
|
||||
std::vector<SparseData<T>> &sparse_data,
|
||||
std::vector<std::vector<uint64_t>> &taglists) {
|
||||
std::fstream qf(query_file, std::ios::in);
|
||||
|
||||
if (!qf.is_open()) {
|
||||
std::cerr << "open query file failed! [" << query_file << "]"
|
||||
<< std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
bool ret;
|
||||
std::string buffer;
|
||||
while (getline(qf, buffer)) {
|
||||
buffer.erase(buffer.find_last_not_of('\n') + 1);
|
||||
if (buffer.empty()) {
|
||||
continue;
|
||||
}
|
||||
std::vector<std::string> res;
|
||||
ailego::StringHelper::Split(buffer, first_sep, &res);
|
||||
if (res.empty()) {
|
||||
continue;
|
||||
}
|
||||
std::string feature_str = res[0];
|
||||
if (res.size() > 1) {
|
||||
feature_str = res[1];
|
||||
}
|
||||
std::vector<T> feature;
|
||||
size_t dimension = 0;
|
||||
ret = load_from_string(feature_str, second_sep, feature, &dimension);
|
||||
if (!ret) {
|
||||
return false;
|
||||
}
|
||||
|
||||
features.emplace_back(feature);
|
||||
|
||||
uint64_t key = atol(res[0].c_str());
|
||||
|
||||
// load sparse feature
|
||||
uint32_t sparse_count = 0;
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
std::vector<T> sparse_feature;
|
||||
|
||||
if (res.size() >= 3) {
|
||||
ret = load_from_string_sparse(key, res[2], second_sep, sparse_indices,
|
||||
sparse_feature, &sparse_count);
|
||||
if (!ret) {
|
||||
std::cerr << "load sparse failed for key: " << key << std::endl;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
sparse_data.emplace_back(sparse_count, std::move(sparse_indices),
|
||||
std::move(sparse_feature));
|
||||
|
||||
if (res.size() >= 4) {
|
||||
std::vector<uint64_t> taglist;
|
||||
size_t tag_count = 0;
|
||||
|
||||
ret = load_tags_from_string(res[4], second_sep, taglist, &tag_count);
|
||||
if (!ret) {
|
||||
std::cerr << "load tags failed for key: " << key << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
taglists.emplace_back(taglist);
|
||||
}
|
||||
}
|
||||
|
||||
qf.close();
|
||||
if (features.size() == 0) {
|
||||
std::cerr << "Read query size is 0" << std::endl;
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
bool load_record(const std::string &input, const std::string &first_sep,
|
||||
const std::string &second_sep, const size_t dimension,
|
||||
std::vector<uint64_t> &keys,
|
||||
std::vector<std::vector<T>> &features,
|
||||
std::vector<SparseData<T>> &sparse_data,
|
||||
std::vector<std::vector<uint64_t>> &taglists) {
|
||||
std::fstream qf(input, std::ios::in);
|
||||
|
||||
if (!qf.is_open()) {
|
||||
std::cerr << "open file failed! [" << input << "]" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
bool ret;
|
||||
uint32_t count = 0;
|
||||
std::string buffer;
|
||||
|
||||
while (getline(qf, buffer)) {
|
||||
buffer.erase(buffer.find_last_not_of('\n') + 1);
|
||||
if (buffer.empty()) {
|
||||
continue;
|
||||
}
|
||||
std::vector<std::string> res;
|
||||
ailego::StringHelper::Split(buffer, first_sep, &res);
|
||||
if (res.size() < 2) {
|
||||
std::cerr << "skip record : " << buffer << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
std::vector<T> feature;
|
||||
size_t real_dim = 0;
|
||||
|
||||
// load sparse feature
|
||||
uint32_t sparse_count = 0;
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
std::vector<T> sparse_feature;
|
||||
|
||||
uint64_t key = atol(res[0].c_str());
|
||||
|
||||
// load dense feature
|
||||
ret = load_from_string(res[1], second_sep, feature, &real_dim);
|
||||
if (!ret) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (real_dim != dimension) {
|
||||
std::cerr << "real dim (" << real_dim << ") is not equal to dimension("
|
||||
<< dimension << ") key : " << res[0] << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
features.emplace_back(feature);
|
||||
keys.emplace_back(key);
|
||||
|
||||
if (res.size() >= 3) {
|
||||
ret = load_from_string_sparse(key, res[2], second_sep, sparse_indices,
|
||||
sparse_feature, &sparse_count);
|
||||
if (!ret) {
|
||||
std::cerr << "load sparse failed for key: " << key << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
sparse_data.emplace_back(sparse_count, std::move(sparse_indices),
|
||||
std::move(sparse_feature));
|
||||
}
|
||||
|
||||
if (res.size() >= 4) {
|
||||
std::vector<uint64_t> taglist;
|
||||
size_t tag_count = 0;
|
||||
|
||||
ret = load_tags_from_string(res[3], second_sep, taglist, &tag_count);
|
||||
if (!ret) {
|
||||
std::cerr << "load tags failed for key: " << key << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
taglists.emplace_back(taglist);
|
||||
}
|
||||
|
||||
count++;
|
||||
if (count % 1000000 == 0) {
|
||||
std::cout << "processed " << count << " records!" << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
qf.close();
|
||||
|
||||
if (keys.size() == 0) {
|
||||
std::cerr << "Reading nothing from input" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_record_sparse(const std::string &input,
|
||||
const std::string &first_sep,
|
||||
const std::string &second_sep,
|
||||
std::vector<uint64_t> &keys,
|
||||
std::vector<SparseData<T>> &sparse_data,
|
||||
std::vector<std::vector<uint64_t>> &taglists) {
|
||||
std::fstream qf(input, std::ios::in);
|
||||
|
||||
if (!qf.is_open()) {
|
||||
std::cerr << "open file failed! [" << input << "]" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
bool ret;
|
||||
uint32_t count = 0;
|
||||
std::string buffer;
|
||||
|
||||
while (getline(qf, buffer)) {
|
||||
buffer.erase(buffer.find_last_not_of('\n') + 1);
|
||||
if (buffer.empty()) {
|
||||
continue;
|
||||
}
|
||||
std::vector<std::string> res;
|
||||
ailego::StringHelper::Split(buffer, first_sep, &res);
|
||||
if (res.size() < 2) {
|
||||
std::cerr << "skip record : " << buffer << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
uint64_t key = atol(res[0].c_str());
|
||||
|
||||
// load sparse feature
|
||||
uint32_t sparse_count = 0;
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
std::vector<T> sparse_feature;
|
||||
|
||||
if (res.size() <= 2) {
|
||||
std::cerr << "field erorr, key: " << key << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
ret = load_from_string_sparse(key, res[2], second_sep, sparse_indices,
|
||||
sparse_feature, &sparse_count);
|
||||
if (!ret) {
|
||||
std::cerr << "load sparse failed for key: " << key << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
keys.emplace_back(key);
|
||||
|
||||
sparse_data.emplace_back(sparse_count, std::move(sparse_indices),
|
||||
std::move(sparse_feature));
|
||||
|
||||
if (res.size() >= 4) {
|
||||
std::vector<uint64_t> taglist;
|
||||
size_t tag_count;
|
||||
|
||||
ret = load_tags_from_string(res[4], second_sep, taglist, &tag_count);
|
||||
if (!ret) {
|
||||
std::cerr << "load tags failed for key: " << key << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
taglists.emplace_back(taglist);
|
||||
}
|
||||
|
||||
count++;
|
||||
if (count % 1000000 == 0) {
|
||||
std::cout << "processed " << count << " records!" << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
qf.close();
|
||||
|
||||
if (keys.size() == 0) {
|
||||
std::cerr << "Reading nothing from input" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
bool load_from_string(const std::string &record,
|
||||
const std::string &second_sep, std::vector<U> &data,
|
||||
size_t *count) {
|
||||
ailego::StringHelper::Split(record, second_sep, &data, true);
|
||||
*count = data.size();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_scores_from_string(const std::string &record,
|
||||
const std::string &second_sep,
|
||||
std::vector<float> &data, size_t *count) {
|
||||
ailego::StringHelper::Split(record, second_sep, &data, true);
|
||||
*count = data.size();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_ids_from_string(const std::string &record,
|
||||
const std::string &second_sep,
|
||||
std::vector<uint64_t> &data, size_t *count) {
|
||||
ailego::StringHelper::Split(record, second_sep, &data, true);
|
||||
*count = data.size();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_tags_from_string(const std::string &record,
|
||||
const std::string &second_sep,
|
||||
std::vector<uint64_t> &tags, size_t *count) {
|
||||
ailego::StringHelper::Split(record, second_sep, &tags, true);
|
||||
*count = tags.size();
|
||||
|
||||
// order tags
|
||||
sort(tags.begin(), tags.end());
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
// overloading for binary
|
||||
bool load_from_string(const std::string &record,
|
||||
const std::string &second_sep,
|
||||
std::vector<uint32_t> &data, size_t *count) {
|
||||
// fetch split value from text file
|
||||
std::vector<uint8_t> vec;
|
||||
ailego::StringHelper::Split(record, second_sep, &vec, true);
|
||||
if (vec.size() == 0) {
|
||||
std::cerr << "Binary vector size is 0" << std::endl;
|
||||
return false;
|
||||
}
|
||||
if (vec.size() % 32 != 0) {
|
||||
std::cerr << "Binary vector size must be 32_X" << std::endl;
|
||||
return false;
|
||||
}
|
||||
// compact into uint32_t
|
||||
size_t sz = vec.size();
|
||||
std::vector<uint8_t> tmp;
|
||||
for (size_t i = 0; i < sz; i += 8) {
|
||||
uint8_t v = 0;
|
||||
v |= (vec[i] & 0x01) << 7;
|
||||
v |= (vec[i + 1] & 0x01) << 6;
|
||||
v |= (vec[i + 2] & 0x01) << 5;
|
||||
v |= (vec[i + 3] & 0x01) << 4;
|
||||
v |= (vec[i + 4] & 0x01) << 3;
|
||||
v |= (vec[i + 5] & 0x01) << 2;
|
||||
v |= (vec[i + 6] & 0x01) << 1;
|
||||
v |= (vec[i + 7] & 0x01) << 0;
|
||||
tmp.push_back(v);
|
||||
}
|
||||
data.resize(sz / 32);
|
||||
memcpy(&data[0], &tmp[0], tmp.size());
|
||||
*count = sz;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
// overloading for binary
|
||||
bool load_from_string(const std::string &record,
|
||||
const std::string &second_sep,
|
||||
std::vector<uint64_t> &data, size_t *count) {
|
||||
// fetch split value from text file
|
||||
std::vector<uint8_t> vec;
|
||||
ailego::StringHelper::Split(record, second_sep, &vec);
|
||||
if (vec.size() == 0) {
|
||||
std::cerr << "Binary vector size is 0" << std::endl;
|
||||
return false;
|
||||
}
|
||||
if (vec.size() % 64 != 0) {
|
||||
std::cerr << "Binary vector size must be 32_X" << std::endl;
|
||||
return false;
|
||||
}
|
||||
// compact into uint64_t
|
||||
size_t sz = vec.size();
|
||||
std::vector<uint8_t> tmp;
|
||||
for (size_t i = 0; i < sz; i += 8) {
|
||||
uint8_t v = 0;
|
||||
v |= (vec[i] & 0x01) << 7;
|
||||
v |= (vec[i + 1] & 0x01) << 6;
|
||||
v |= (vec[i + 2] & 0x01) << 5;
|
||||
v |= (vec[i + 3] & 0x01) << 4;
|
||||
v |= (vec[i + 4] & 0x01) << 3;
|
||||
v |= (vec[i + 5] & 0x01) << 2;
|
||||
v |= (vec[i + 6] & 0x01) << 1;
|
||||
v |= (vec[i + 7] & 0x01) << 0;
|
||||
tmp.push_back(v);
|
||||
}
|
||||
data.resize(sz / 64);
|
||||
memcpy(&data[0], &tmp[0], tmp.size());
|
||||
*count = sz;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_from_string_sparse(uint64_t key, const std::string &record,
|
||||
const std::string &second_sep,
|
||||
std::vector<uint32_t> &sparse_indices,
|
||||
std::vector<T> &sparse_feature,
|
||||
uint32_t *sparse_count) {
|
||||
const std::string sparse_sep = ":";
|
||||
std::vector<std::string> res;
|
||||
ailego::StringHelper::Split(record, sparse_sep, &res);
|
||||
|
||||
if (res.size() == 2) {
|
||||
ailego::StringHelper::Split(res[0], second_sep, &sparse_indices);
|
||||
ailego::StringHelper::Split(res[1], second_sep, &sparse_feature);
|
||||
|
||||
uint32_t index_count = sparse_indices.size();
|
||||
uint32_t feature_count = sparse_feature.size();
|
||||
|
||||
if (feature_count == index_count) {
|
||||
*sparse_count = feature_count;
|
||||
} else {
|
||||
std::cerr << "sparse feature count (" << feature_count
|
||||
<< ") is not equal with sparse index count(" << index_count
|
||||
<< ") key : " << key << std::endl;
|
||||
*sparse_count = 0;
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
// check order
|
||||
for (size_t i = 1; i < sparse_indices.size(); ++i) {
|
||||
if (sparse_indices[i - 1] >= sparse_indices[i]) {
|
||||
std::cerr << "sparse indices not ordered, key : " << key
|
||||
<< ", dim info: [" << sparse_indices[i - 1] << ", "
|
||||
<< sparse_indices[i] << "]" << std::endl;
|
||||
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
// LINE FORMAT is as follows:
|
||||
// key:key0 key1 key2 ... keyN:score0 score1 score2 ... scoreN
|
||||
bool load_external_gt(
|
||||
const std::string &input, const std::string &first_sep,
|
||||
const std::string &second_sep,
|
||||
std::vector<std::vector<std::pair<uint64_t, float>>> &ground_truth) {
|
||||
std::fstream gf(input, std::ios::in);
|
||||
|
||||
if (!gf.is_open()) {
|
||||
std::cerr << "open file failed! [" << input << "]" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
uint32_t count = 0;
|
||||
std::string buffer;
|
||||
while (getline(gf, buffer)) {
|
||||
buffer.erase(buffer.find_last_not_of('\n') + 1);
|
||||
if (buffer.empty()) {
|
||||
continue;
|
||||
}
|
||||
std::vector<std::string> res;
|
||||
ailego::StringHelper::Split(buffer, first_sep, &res);
|
||||
if (res.size() < 2) {
|
||||
std::cerr << "skip record : " << buffer << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
// uint64_t main_key = std::strtoll(res[0].c_str(), NULL, 10);
|
||||
if (res.size() == 2) {
|
||||
std::vector<uint64_t> keys;
|
||||
size_t key_num = 0;
|
||||
load_ids_from_string(res[1], second_sep, keys, &key_num);
|
||||
|
||||
std::vector<std::pair<uint64_t, float>> one_groud_truth;
|
||||
for (size_t i = 0; i < keys.size(); ++i) {
|
||||
one_groud_truth.push_back(std::make_pair(keys[i], 0.0f));
|
||||
}
|
||||
|
||||
ground_truth.push_back(std::move(one_groud_truth));
|
||||
} else {
|
||||
std::vector<uint64_t> keys;
|
||||
size_t key_num = 0;
|
||||
load_ids_from_string(res[1], second_sep, keys, &key_num);
|
||||
|
||||
std::vector<float> scores;
|
||||
size_t score_num = 0;
|
||||
load_scores_from_string(res[2], second_sep, scores, &score_num);
|
||||
|
||||
if (key_num != score_num) {
|
||||
std::cerr << "key num (" << key_num << ") is not equal to ("
|
||||
<< score_num << "), line data:" << buffer << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
std::vector<std::pair<uint64_t, float>> one_groud_truth;
|
||||
for (size_t i = 0; i < keys.size(); ++i) {
|
||||
one_groud_truth.push_back(std::make_pair(keys[i], scores[i]));
|
||||
}
|
||||
|
||||
ground_truth.push_back(std::move(one_groud_truth));
|
||||
}
|
||||
|
||||
count++;
|
||||
if (count % 1000000 == 0) {
|
||||
std::cout << "processed " << count << " records!" << std::endl;
|
||||
}
|
||||
}
|
||||
gf.close();
|
||||
if (ground_truth.size() == 0) {
|
||||
std::cerr << "Reading nothing from input" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
@@ -0,0 +1,56 @@
|
||||
// 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 <cstdint>
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
enum VecsBitMapIndex {
|
||||
BITMAP_INDEX_KEY = 0,
|
||||
BITMAP_INDEX_DENSE = 1,
|
||||
BITMAP_INDEX_SPARSE = 2,
|
||||
BITMAP_INDEX_TAGLIST = 4
|
||||
};
|
||||
|
||||
#pragma pack(4)
|
||||
struct VecsHeader {
|
||||
uint64_t num_vecs;
|
||||
uint16_t meta_size_v1;
|
||||
uint16_t version;
|
||||
uint32_t meta_size;
|
||||
uint64_t bitmap; // set for data section
|
||||
uint64_t key_offset; // offset for key
|
||||
uint64_t key_size; // size for key
|
||||
uint64_t dense_offset; // offset for dense
|
||||
uint64_t dense_size; // size for dense
|
||||
uint64_t sparse_offset; // offset for sparse
|
||||
uint64_t sparse_size; // size for sparse
|
||||
uint64_t partition_offset; // offset for partition
|
||||
uint64_t partition_size; // size for partition
|
||||
uint64_t taglist_offset; // offset for taglist
|
||||
uint64_t taglist_size; // size for taglist
|
||||
|
||||
uint8_t *meta_buf() {
|
||||
return reinterpret_cast<uint8_t *>(this + 1);
|
||||
}
|
||||
const uint8_t *meta_buf() const {
|
||||
return reinterpret_cast<const uint8_t *>(this + 1);
|
||||
}
|
||||
};
|
||||
#pragma pack()
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
@@ -0,0 +1,433 @@
|
||||
// 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 <string>
|
||||
#include <unordered_map>
|
||||
#include <zvec/ailego/container/params.h>
|
||||
#include "zvec/core/framework/index_error.h"
|
||||
#include "zvec/core/framework/index_holder.h"
|
||||
#include "zvec/core/framework/index_provider.h"
|
||||
#include "zvec/core/framework/index_storage.h"
|
||||
#include "vecs_reader.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
/*!
|
||||
* Vecs Index Holder
|
||||
* framwork will use IndexHolder in this way:
|
||||
* for (iter = create_iterator(); iter->is_valid(); iter->next()) {
|
||||
* key = iter->key();
|
||||
* data = iter->data();
|
||||
* }
|
||||
*/
|
||||
class VecsIndexHolder : public IndexProvider {
|
||||
public:
|
||||
typedef std::shared_ptr<VecsIndexHolder> Pointer;
|
||||
|
||||
bool load(const std::string &file_path) {
|
||||
if (!vecs_reader_.load(file_path)) {
|
||||
return false;
|
||||
}
|
||||
build_key_index_map();
|
||||
return true;
|
||||
}
|
||||
|
||||
const IndexMeta &index_meta(void) const {
|
||||
return vecs_reader_.index_meta();
|
||||
}
|
||||
|
||||
void set_metric(const std::string &name, const ailego::Params ¶ms) {
|
||||
vecs_reader_.set_metric(name, params);
|
||||
}
|
||||
|
||||
/*!
|
||||
* Index Holder Iterator
|
||||
*/
|
||||
class Iterator : public IndexHybridHolder::Iterator {
|
||||
public:
|
||||
//! Constructor
|
||||
Iterator(const VecsIndexHolder &holder, uint32_t cursor)
|
||||
: cursor_(cursor),
|
||||
vecs_reader_(holder.vecs_reader_),
|
||||
stop_(holder.stop_) {}
|
||||
|
||||
//! Test if the iterator is valid
|
||||
virtual bool is_valid(void) const override {
|
||||
return !stop_ && cursor_ < vecs_reader_.num_vecs();
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual uint64_t key(void) const override {
|
||||
return vecs_reader_.get_key(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve pointer of data
|
||||
virtual const void *data() const override {
|
||||
return vecs_reader_.get_vector(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual uint32_t sparse_count() const override {
|
||||
return vecs_reader_.get_sparse_count(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual const uint32_t *sparse_indices() const override {
|
||||
return vecs_reader_.get_sparse_indices(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual const void *sparse_data() const override {
|
||||
return vecs_reader_.get_sparse_data(cursor_);
|
||||
}
|
||||
|
||||
//! Next iterator
|
||||
virtual void next(void) override {
|
||||
++cursor_;
|
||||
}
|
||||
|
||||
//! Reset the iterator
|
||||
virtual void reset(void) {
|
||||
cursor_ = 0;
|
||||
}
|
||||
|
||||
private:
|
||||
size_t cursor_;
|
||||
const VecsReader &vecs_reader_;
|
||||
const bool &stop_;
|
||||
};
|
||||
|
||||
virtual IndexHolder::Iterator::Pointer create_iterator(void) override {
|
||||
// make sure iter has value whenn create_iterator finished
|
||||
IndexHolder::Iterator::Pointer iter(
|
||||
new VecsIndexHolder::Iterator(*this, start_cursor_));
|
||||
return iter;
|
||||
}
|
||||
|
||||
virtual IndexHybridHolder::Iterator::Pointer create_hybrid_iterator(void) {
|
||||
// make sure iter has value whenn create_iterator finished
|
||||
IndexHybridHolder::Iterator::Pointer iter(
|
||||
new VecsIndexHolder::Iterator(*this, start_cursor_));
|
||||
return iter;
|
||||
}
|
||||
|
||||
//! Retrieve count of elements in holder
|
||||
virtual size_t count(void) const override {
|
||||
return max_doc_count_ != 0
|
||||
? std::min(max_doc_count_, vecs_reader_.num_vecs())
|
||||
: vecs_reader_.num_vecs();
|
||||
}
|
||||
|
||||
//! Retrieve dimension
|
||||
virtual size_t dimension(void) const override {
|
||||
return vecs_reader_.index_meta().dimension();
|
||||
}
|
||||
|
||||
//! Retrieve type information
|
||||
virtual IndexMeta::DataType data_type(void) const override {
|
||||
return vecs_reader_.index_meta().data_type();
|
||||
}
|
||||
|
||||
//! Retrieve element size in bytes
|
||||
virtual size_t element_size(void) const override {
|
||||
return vecs_reader_.index_meta().element_size();
|
||||
}
|
||||
|
||||
//! Retrieve if it can multi-pass
|
||||
virtual bool multipass(void) const override {
|
||||
return true;
|
||||
}
|
||||
|
||||
void stop(void) {
|
||||
stop_ = true;
|
||||
}
|
||||
|
||||
uint64_t get_num_vecs() const {
|
||||
return vecs_reader_.num_vecs();
|
||||
}
|
||||
|
||||
uint64_t get_key(size_t idx) const {
|
||||
return vecs_reader_.get_key(idx);
|
||||
}
|
||||
|
||||
uint32_t get_sparse_count(size_t idx) const {
|
||||
return vecs_reader_.get_sparse_count(idx);
|
||||
}
|
||||
|
||||
const uint32_t *get_sparse_indices(size_t idx) const {
|
||||
return vecs_reader_.get_sparse_indices(idx);
|
||||
}
|
||||
|
||||
const void *get_sparse_data(size_t idx) const {
|
||||
return vecs_reader_.get_sparse_data(idx);
|
||||
}
|
||||
|
||||
void set_start_cursor(uint32_t index) {
|
||||
start_cursor_ = index;
|
||||
}
|
||||
|
||||
void set_max_doc_count(size_t value) {
|
||||
max_doc_count_ = value;
|
||||
}
|
||||
|
||||
uint32_t start_cursor() const {
|
||||
return start_cursor_;
|
||||
}
|
||||
|
||||
size_t total_sparse_count(void) const {
|
||||
return vecs_reader_.get_total_sparse_count();
|
||||
}
|
||||
|
||||
bool has_taglist() const {
|
||||
return vecs_reader_.has_taglist();
|
||||
}
|
||||
|
||||
uint64_t get_taglist_count(size_t index) const {
|
||||
return vecs_reader_.get_taglist_count(index);
|
||||
}
|
||||
|
||||
const void *get_taglist(size_t index) const {
|
||||
return vecs_reader_.get_taglist(index);
|
||||
}
|
||||
|
||||
const void *get_taglist_data(size_t &size) const {
|
||||
return vecs_reader_.get_taglist_data(size);
|
||||
}
|
||||
|
||||
const void *get_key_base() const {
|
||||
return vecs_reader_.key_base();
|
||||
}
|
||||
|
||||
const void *get_vector_by_index(size_t idx) const {
|
||||
return vecs_reader_.get_vector(idx);
|
||||
}
|
||||
|
||||
public: // IndexProvider interface implementation
|
||||
//! Retrieve a vector using a primary key
|
||||
const void *get_vector(const uint64_t key) const override {
|
||||
auto it = key_to_index_map_.find(key);
|
||||
if (it == key_to_index_map_.end()) {
|
||||
return nullptr;
|
||||
}
|
||||
return vecs_reader_.get_vector(it->second);
|
||||
}
|
||||
|
||||
//! Retrieve a vector using a primary key
|
||||
virtual int get_vector(const uint64_t key,
|
||||
IndexStorage::MemoryBlock &block) const override {
|
||||
const void *vector = get_vector(key);
|
||||
if (vector == nullptr) {
|
||||
return IndexError_NoExist;
|
||||
}
|
||||
block.reset((void *)vector);
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Retrieve the owner class
|
||||
virtual const std::string &owner_class(void) const override {
|
||||
static std::string owner_class_name = "VecsIndexHolder";
|
||||
return owner_class_name;
|
||||
}
|
||||
|
||||
private:
|
||||
//! Build key to index mapping
|
||||
void build_key_index_map() {
|
||||
key_to_index_map_.clear();
|
||||
size_t num_vecs = vecs_reader_.num_vecs();
|
||||
for (size_t i = 0; i < num_vecs; ++i) {
|
||||
uint64_t key = vecs_reader_.get_key(i);
|
||||
key_to_index_map_[key] = i;
|
||||
}
|
||||
}
|
||||
|
||||
bool stop_{false};
|
||||
uint32_t start_cursor_{0};
|
||||
VecsReader vecs_reader_;
|
||||
size_t max_doc_count_{0};
|
||||
std::unordered_map<uint64_t, size_t> key_to_index_map_;
|
||||
};
|
||||
|
||||
|
||||
/*!
|
||||
* Vecs Index Sparse Holder
|
||||
* framwork will use IndexHolder in this way:
|
||||
* for (iter = create_iterator(); iter->is_valid(); iter->next()) {
|
||||
* key = iter->key();
|
||||
* data = iter->sparse_data();
|
||||
* }
|
||||
*/
|
||||
class VecsIndexSparseHolder : public IndexSparseHolder {
|
||||
public:
|
||||
typedef std::shared_ptr<VecsIndexSparseHolder> Pointer;
|
||||
|
||||
bool load(const std::string &file_path) {
|
||||
return vecs_reader_.load(file_path);
|
||||
}
|
||||
|
||||
const IndexMeta &index_meta(void) const {
|
||||
return vecs_reader_.index_meta();
|
||||
}
|
||||
|
||||
void set_metric(const std::string &name, const ailego::Params ¶ms) {
|
||||
vecs_reader_.set_metric(name, params);
|
||||
}
|
||||
|
||||
/*!
|
||||
* Index Holder Iterator
|
||||
*/
|
||||
class Iterator : public IndexSparseHolder::Iterator {
|
||||
public:
|
||||
//! Constructor
|
||||
Iterator(const VecsIndexSparseHolder &holder, uint32_t cursor)
|
||||
: cursor_(cursor),
|
||||
vecs_reader_(holder.vecs_reader_),
|
||||
stop_(holder.stop_) {}
|
||||
|
||||
//! Test if the iterator is valid
|
||||
virtual bool is_valid(void) const override {
|
||||
return !stop_ && cursor_ < vecs_reader_.num_vecs();
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual uint64_t key(void) const override {
|
||||
return vecs_reader_.get_key(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual uint32_t sparse_count() const override {
|
||||
return vecs_reader_.get_sparse_count(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual const uint32_t *sparse_indices() const override {
|
||||
return vecs_reader_.get_sparse_indices(cursor_);
|
||||
}
|
||||
|
||||
//! Retrieve primary key
|
||||
virtual const void *sparse_data() const override {
|
||||
return vecs_reader_.get_sparse_data(cursor_);
|
||||
}
|
||||
|
||||
//! Next iterator
|
||||
virtual void next(void) override {
|
||||
++cursor_;
|
||||
}
|
||||
|
||||
//! Reset the iterator
|
||||
virtual void reset(void) {
|
||||
cursor_ = 0;
|
||||
}
|
||||
|
||||
private:
|
||||
size_t cursor_;
|
||||
const SparseVecsReader &vecs_reader_;
|
||||
const bool &stop_;
|
||||
};
|
||||
|
||||
virtual IndexSparseHolder::Iterator::Pointer create_iterator(void) override {
|
||||
// make sure iter has value whenn create_iterator finished
|
||||
IndexSparseHolder::Iterator::Pointer iter(
|
||||
new VecsIndexSparseHolder::Iterator(*this, start_cursor_));
|
||||
return iter;
|
||||
}
|
||||
|
||||
//! Retrieve count of elements in holder
|
||||
virtual size_t count(void) const override {
|
||||
return max_doc_count_ != 0
|
||||
? std::min(max_doc_count_, vecs_reader_.num_vecs())
|
||||
: vecs_reader_.num_vecs();
|
||||
}
|
||||
|
||||
//! Retrieve type information
|
||||
virtual IndexMeta::DataType data_type(void) const override {
|
||||
return vecs_reader_.index_meta().data_type();
|
||||
}
|
||||
|
||||
//! Retrieve if it can multi-pass
|
||||
virtual bool multipass(void) const override {
|
||||
return true;
|
||||
}
|
||||
|
||||
void stop(void) {
|
||||
stop_ = true;
|
||||
}
|
||||
|
||||
uint64_t get_key(size_t idx) const {
|
||||
return vecs_reader_.get_key(idx);
|
||||
}
|
||||
|
||||
uint32_t get_sparse_count(size_t idx) const {
|
||||
return vecs_reader_.get_sparse_count(idx);
|
||||
}
|
||||
|
||||
const uint32_t *get_sparse_indices(size_t idx) const {
|
||||
return vecs_reader_.get_sparse_indices(idx);
|
||||
}
|
||||
|
||||
const void *get_sparse_data(size_t idx) const {
|
||||
return vecs_reader_.get_sparse_data(idx);
|
||||
}
|
||||
|
||||
void set_start_cursor(uint32_t index) {
|
||||
start_cursor_ = index;
|
||||
}
|
||||
|
||||
void set_max_doc_count(size_t value) {
|
||||
max_doc_count_ = value;
|
||||
}
|
||||
|
||||
uint64_t get_num_vecs() const {
|
||||
return vecs_reader_.num_vecs();
|
||||
}
|
||||
|
||||
uint32_t start_cursor() const {
|
||||
return start_cursor_;
|
||||
}
|
||||
|
||||
size_t total_sparse_count(void) const override {
|
||||
return vecs_reader_.get_total_sparse_count();
|
||||
}
|
||||
|
||||
bool has_taglist() const {
|
||||
return vecs_reader_.has_taglist();
|
||||
}
|
||||
|
||||
uint64_t get_taglist_count(size_t index) const {
|
||||
return vecs_reader_.get_taglist_count(index);
|
||||
}
|
||||
|
||||
const void *get_taglist(size_t index) const {
|
||||
return vecs_reader_.get_taglist(index);
|
||||
}
|
||||
|
||||
const void *get_taglist_data(size_t &size) const {
|
||||
return vecs_reader_.get_taglist_data(size);
|
||||
}
|
||||
|
||||
const void *get_key_base() const {
|
||||
return vecs_reader_.key_base();
|
||||
}
|
||||
|
||||
private:
|
||||
bool stop_{false};
|
||||
uint32_t start_cursor_{0};
|
||||
SparseVecsReader vecs_reader_;
|
||||
size_t max_doc_count_{0};
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
@@ -0,0 +1,378 @@
|
||||
// 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
|
||||
Reference in New Issue
Block a user