wehub-resource-sync
264 lines
8.9 KiB
C++
264 lines
8.9 KiB
C++
#include <cstdlib>
|
|
#include <filesystem>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <zvec/db/collection.h>
|
|
#include <zvec/db/doc.h>
|
|
#include <zvec/db/schema.h>
|
|
#include <zvec/db/status.h>
|
|
|
|
using namespace zvec;
|
|
|
|
Doc create_doc(const uint64_t doc_id, const CollectionSchema &schema,
|
|
std::string pk = "") {
|
|
Doc new_doc;
|
|
if (pk.empty()) {
|
|
pk = "pk_" + std::to_string(doc_id);
|
|
}
|
|
new_doc.set_pk(pk);
|
|
|
|
for (auto &field : schema.fields()) {
|
|
switch (field->data_type()) {
|
|
case DataType::BINARY: {
|
|
std::string binary_str("binary_" + std::to_string(doc_id));
|
|
new_doc.set<std::string>(field->name(), binary_str);
|
|
break;
|
|
}
|
|
case DataType::BOOL:
|
|
new_doc.set<bool>(field->name(), doc_id % 10 == 0);
|
|
break;
|
|
case DataType::INT32:
|
|
new_doc.set<int32_t>(field->name(), (int32_t)doc_id);
|
|
break;
|
|
case DataType::INT64:
|
|
new_doc.set<int64_t>(field->name(), (int64_t)doc_id);
|
|
break;
|
|
case DataType::UINT32:
|
|
new_doc.set<uint32_t>(field->name(), (uint32_t)doc_id);
|
|
break;
|
|
case DataType::UINT64:
|
|
new_doc.set<uint64_t>(field->name(), (uint64_t)doc_id);
|
|
break;
|
|
case DataType::FLOAT:
|
|
new_doc.set<float>(field->name(), (float)doc_id);
|
|
break;
|
|
case DataType::DOUBLE:
|
|
new_doc.set<double>(field->name(), (double)doc_id);
|
|
break;
|
|
case DataType::STRING:
|
|
new_doc.set<std::string>(field->name(),
|
|
"value_" + std::to_string(doc_id));
|
|
break;
|
|
case DataType::ARRAY_BINARY: {
|
|
std::vector<std::string> bin_vec;
|
|
for (size_t i = 0; i < (doc_id % 10); i++) {
|
|
bin_vec.push_back("bin_" + std::to_string(i));
|
|
}
|
|
new_doc.set<std::vector<std::string>>(field->name(), bin_vec);
|
|
break;
|
|
}
|
|
case DataType::ARRAY_BOOL:
|
|
new_doc.set<std::vector<bool>>(field->name(),
|
|
std::vector<bool>(10, doc_id % 10 == 0));
|
|
break;
|
|
case DataType::ARRAY_INT32:
|
|
new_doc.set<std::vector<int32_t>>(
|
|
field->name(), std::vector<int32_t>(10, (int32_t)doc_id));
|
|
break;
|
|
case DataType::ARRAY_INT64:
|
|
new_doc.set<std::vector<int64_t>>(
|
|
field->name(), std::vector<int64_t>(10, (int64_t)doc_id));
|
|
break;
|
|
case DataType::ARRAY_UINT32:
|
|
new_doc.set<std::vector<uint32_t>>(
|
|
field->name(), std::vector<uint32_t>(10, (uint32_t)doc_id));
|
|
break;
|
|
case DataType::ARRAY_UINT64:
|
|
new_doc.set<std::vector<uint64_t>>(
|
|
field->name(), std::vector<uint64_t>(10, (uint64_t)doc_id));
|
|
break;
|
|
case DataType::ARRAY_FLOAT:
|
|
new_doc.set<std::vector<float>>(field->name(),
|
|
std::vector<float>(10, (float)doc_id));
|
|
break;
|
|
case DataType::ARRAY_DOUBLE:
|
|
new_doc.set<std::vector<double>>(
|
|
field->name(), std::vector<double>(10, (double)doc_id));
|
|
break;
|
|
case DataType::ARRAY_STRING:
|
|
new_doc.set<std::vector<std::string>>(
|
|
field->name(),
|
|
std::vector<std::string>(10, "value_" + std::to_string(doc_id)));
|
|
break;
|
|
case DataType::VECTOR_BINARY32:
|
|
new_doc.set<std::vector<uint32_t>>(
|
|
field->name(),
|
|
std::vector<uint32_t>(field->dimension(), uint32_t(doc_id + 0.1)));
|
|
break;
|
|
case DataType::VECTOR_BINARY64:
|
|
new_doc.set<std::vector<uint64_t>>(
|
|
field->name(),
|
|
std::vector<uint64_t>(field->dimension(), uint64_t(doc_id + 0.1)));
|
|
break;
|
|
case DataType::VECTOR_FP32:
|
|
new_doc.set<std::vector<float>>(
|
|
field->name(),
|
|
std::vector<float>(field->dimension(), float(doc_id + 0.1)));
|
|
break;
|
|
case DataType::VECTOR_FP64:
|
|
new_doc.set<std::vector<double>>(
|
|
field->name(),
|
|
std::vector<double>(field->dimension(), double(doc_id + 0.1)));
|
|
break;
|
|
case DataType::VECTOR_FP16:
|
|
new_doc.set<std::vector<zvec::float16_t>>(
|
|
field->name(), std::vector<zvec::float16_t>(
|
|
field->dimension(), static_cast<zvec::float16_t>(
|
|
float(doc_id + 0.1))));
|
|
break;
|
|
case DataType::VECTOR_INT8:
|
|
new_doc.set<std::vector<int8_t>>(
|
|
field->name(),
|
|
std::vector<int8_t>(field->dimension(), (int8_t)doc_id));
|
|
break;
|
|
case DataType::VECTOR_INT16:
|
|
new_doc.set<std::vector<int16_t>>(
|
|
field->name(),
|
|
std::vector<int16_t>(field->dimension(), (int16_t)doc_id));
|
|
break;
|
|
case DataType::SPARSE_VECTOR_FP16: {
|
|
std::vector<uint32_t> indices;
|
|
std::vector<zvec::float16_t> values;
|
|
for (uint32_t i = 0; i < 100; i++) {
|
|
indices.push_back(i);
|
|
values.push_back(zvec::float16_t(float(doc_id + 0.1)));
|
|
}
|
|
std::pair<std::vector<uint32_t>, std::vector<zvec::float16_t>>
|
|
sparse_float_vec;
|
|
sparse_float_vec.first = indices;
|
|
sparse_float_vec.second = values;
|
|
new_doc.set<
|
|
std::pair<std::vector<uint32_t>, std::vector<zvec::float16_t>>>(
|
|
field->name(), sparse_float_vec);
|
|
break;
|
|
}
|
|
case DataType::SPARSE_VECTOR_FP32: {
|
|
std::vector<uint32_t> indices;
|
|
std::vector<float> values;
|
|
for (uint32_t i = 0; i < 100; i++) {
|
|
indices.push_back(i);
|
|
values.push_back(float(doc_id + 0.1));
|
|
}
|
|
std::pair<std::vector<uint32_t>, std::vector<float>> sparse_float_vec;
|
|
sparse_float_vec.first = indices;
|
|
sparse_float_vec.second = values;
|
|
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
|
|
field->name(), sparse_float_vec);
|
|
break;
|
|
}
|
|
default:
|
|
std::cout << "Unsupported data type: " << field->name() << std::endl;
|
|
throw std::runtime_error("Unsupported vector data type");
|
|
}
|
|
}
|
|
|
|
return new_doc;
|
|
}
|
|
|
|
CollectionSchema::Ptr create_schema() {
|
|
auto schema = std::make_shared<CollectionSchema>("demo");
|
|
schema->set_max_doc_count_per_segment(1000);
|
|
|
|
schema->add_field(std::make_shared<FieldSchema>(
|
|
"id", DataType::INT64, false, std::make_shared<InvertIndexParams>(true)));
|
|
schema->add_field(std::make_shared<FieldSchema>(
|
|
"name", DataType::STRING, false,
|
|
std::make_shared<InvertIndexParams>(false)));
|
|
schema->add_field(
|
|
std::make_shared<FieldSchema>("weight", DataType::FLOAT, true));
|
|
|
|
schema->add_field(std::make_shared<FieldSchema>(
|
|
"dense", DataType::VECTOR_FP32, 128, false,
|
|
std::make_shared<HnswIndexParams>(MetricType::IP)));
|
|
schema->add_field(std::make_shared<FieldSchema>(
|
|
"sparse", DataType::SPARSE_VECTOR_FP32, 0, false,
|
|
std::make_shared<HnswIndexParams>(MetricType::IP)));
|
|
|
|
return schema;
|
|
}
|
|
|
|
int main() {
|
|
std::string path = "./demo";
|
|
std::filesystem::remove_all(path);
|
|
|
|
auto schema = create_schema();
|
|
CollectionOptions options{false, true};
|
|
|
|
auto result = Collection::CreateAndOpen(path, *schema, options);
|
|
if (!result.has_value()) {
|
|
std::cout << result.error().message() << std::endl;
|
|
return -1;
|
|
}
|
|
|
|
std::cout << "init stats: " << result.value()->Stats().value().to_string()
|
|
<< std::endl;
|
|
|
|
auto coll = std::move(result).value();
|
|
|
|
// insert docs
|
|
{
|
|
auto doc1 = create_doc(0, *schema);
|
|
std::vector<Doc> docs{doc1};
|
|
auto res = coll->Insert(docs);
|
|
if (!res.has_value()) {
|
|
std::cout << res.error().message() << std::endl;
|
|
return -1;
|
|
}
|
|
std::cout << "after insert stats " << coll->Stats().value().to_string()
|
|
<< std::endl;
|
|
}
|
|
|
|
// optimize
|
|
{
|
|
auto res = coll->Optimize();
|
|
if (!res.ok()) {
|
|
std::cout << res.message() << std::endl;
|
|
return -1;
|
|
}
|
|
std::cout << "after optimize stats " << coll->Stats().value().to_string()
|
|
<< std::endl;
|
|
}
|
|
|
|
// query
|
|
{
|
|
SearchQuery query;
|
|
query.topk_ = 10;
|
|
query.target_.field_name_ = "dense";
|
|
query.include_vector_ = true;
|
|
std::vector<float> query_vector = std::vector<float>(128, 0.1);
|
|
query.target_.set_vector(std::string((char *)query_vector.data(),
|
|
query_vector.size() * sizeof(float)));
|
|
auto res = coll->Query(query);
|
|
if (!res.has_value()) {
|
|
std::cout << res.error().message() << std::endl;
|
|
return -1;
|
|
}
|
|
std::cout << "query result: doc_count[" << res.value().size() << "]"
|
|
<< std::endl;
|
|
std::cout << "first doc: " << res.value()[0]->to_detail_string()
|
|
<< std::endl;
|
|
}
|
|
|
|
// close and reopen
|
|
coll.reset();
|
|
options.read_only_ = true;
|
|
result = Collection::Open(path, options);
|
|
if (!result.has_value()) {
|
|
std::cout << result.error().message() << std::endl;
|
|
return -1;
|
|
}
|
|
std::cout << "reopen stats: " << result.value()->Stats().value().to_string()
|
|
<< std::endl;
|
|
|
|
return 0;
|
|
} |