wehub-resource-sync
88 lines
2.5 KiB
C++
88 lines
2.5 KiB
C++
#include <cstdlib>
|
|
#include <filesystem>
|
|
#include <iostream>
|
|
#include <zvec/core/interface/index.h>
|
|
#include <zvec/core/interface/index_factory.h>
|
|
#include <zvec/core/interface/index_param.h>
|
|
#include <zvec/core/interface/index_param_builders.h>
|
|
|
|
using namespace zvec::core_interface;
|
|
|
|
constexpr uint32_t kDimension = 64;
|
|
const std::string index_name{"test.index"};
|
|
|
|
Index::Pointer create_index(const BaseIndexParam::Pointer ¶m,
|
|
int doc_num = 10) {
|
|
auto index = IndexFactory::CreateAndInitIndex(*param);
|
|
if (!index) {
|
|
std::cout << "Failed to create index." << std::endl;
|
|
return nullptr;
|
|
}
|
|
|
|
int ret = index->Open(
|
|
index_name, StorageOptions{StorageOptions::StorageType::kMMAP, true});
|
|
if (ret != 0) {
|
|
std::cout << "Failed to open index." << std::endl;
|
|
return nullptr;
|
|
}
|
|
|
|
for (int i = 0; i < doc_num; ++i) {
|
|
std::vector<float> vector(kDimension, i / 10.0f + 0.1f);
|
|
VectorData vector_data;
|
|
vector_data.vector = DenseVector{vector.data()};
|
|
ret = index->Add(vector_data, i);
|
|
if (ret != 0) {
|
|
std::cout << "Failed to add to index." << std::endl;
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
ret = index->Train();
|
|
if (ret != 0) {
|
|
std::cout << "Failed to train index." << std::endl;
|
|
return nullptr;
|
|
}
|
|
|
|
return index;
|
|
}
|
|
|
|
int main() {
|
|
std::filesystem::remove(index_name);
|
|
|
|
auto param = HNSWIndexParamBuilder()
|
|
.WithMetricType(MetricType::kInnerProduct)
|
|
.WithDataType(DataType::DT_FP32)
|
|
.WithDimension(kDimension)
|
|
.WithIsSparse(false)
|
|
.Build();
|
|
auto index = create_index(param, 1);
|
|
std::cout << "index stats: " << index->GetDocCount() << std::endl;
|
|
|
|
// query
|
|
auto query_param = HNSWQueryParamBuilder()
|
|
.with_topk(10)
|
|
.with_fetch_vector(true)
|
|
.with_ef_search(20)
|
|
.build();
|
|
|
|
SearchResult result;
|
|
VectorData query;
|
|
std::vector<float> vector(kDimension, 0.1f);
|
|
query.vector = DenseVector{vector.data()};
|
|
int ret = index->Search(query, query_param, &result);
|
|
if (ret != 0) {
|
|
std::cout << "Failed to search index." << std::endl;
|
|
return -1;
|
|
}
|
|
|
|
std::cout << "query results: " << result.doc_list_.size() << std::endl;
|
|
if (result.doc_list_.size() == 0) {
|
|
std::cout << "No results found." << std::endl;
|
|
return -1;
|
|
}
|
|
|
|
std::cout << "key: " << result.doc_list_[0].key()
|
|
<< ", score: " << result.doc_list_[0].score() << std::endl;
|
|
|
|
return 0;
|
|
} |