diff --git a/CHANGELOG.md b/CHANGELOG.md index b6c1e323ec..5bc40d77b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,11 @@ Please mark all change in change log and use the ticket from JIRA. - \#340 - Test cases run failed on 0.6.0 - \#353 - Rename config.h.in to version.h.in - \#374 - sdk_simple return empty result +- \#377 - Create partition success if tag name only contains spaces - \#397 - sdk_simple return incorrect result +- \#399 - Create partition should be failed if partition tag existed +- \#412 - Message returned is confused when partition created with null partition name +- \#416 - Drop the same partition success repeatally - \#440 - Query API in customization still uses old version ## Feature @@ -20,18 +24,20 @@ Please mark all change in change log and use the ticket from JIRA. - \#77 - Support table partition - \#127 - Support new Index type IVFPQ - \#226 - Experimental shards middleware for Milvus +- \#227 - Support new index types SPTAG-KDT and SPTAG-BKT - \#346 - Support build index with multiple gpu ## Improvement -- \#275 - Rename C++ SDK IndexType -- \#284 - Change C++ SDK to shared library +- \#255 - Add ivfsq8 test report detailed version - \#260 - C++ SDK README - \#266 - Rpc request source code refactor -- \#314 - add Find FAISS in CMake +- \#275 - Rename C++ SDK IndexType +- \#284 - Change C++ SDK to shared library +- \#306 - Use int64 for all config integer - \#310 - Add Q&A for 'protocol https not supported or disable in libcurl' issue +- \#314 - add Find FAISS in CMake - \#322 - Add option to enable / disable prometheus - \#358 - Add more information in build.sh and install.md -- \#255 - Add ivfsq8 test report detailed version - \#404 - Add virtual method Init() in Pass abstract class - \#409 - Add a Fallback pass in optimizer diff --git a/README.md b/README.md index 311d0d6843..4a1151dbf8 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ Milvus is the world's fastest similarity search engine for massive-scale feature For more detailed introduction of Milvus and its architecture, see [Milvus overview](https://www.milvus.io/docs/en/aboutmilvus/overview/). -Milvus provides stable [Python](https://github.com/milvus-io/pymilvus), [Java](https://github.com/milvus-io/milvus-sdk-java) and C++ APIs. +Milvus provides stable [Python](https://github.com/milvus-io/pymilvus), [Java](https://github.com/milvus-io/milvus-sdk-java) and [C++](https://github.com/milvus-io/milvus/tree/master/core/src/sdk) APIs. Keep up-to-date with newest releases and latest updates by reading Milvus [release notes](https://www.milvus.io/docs/en/release/v0.5.0/). diff --git a/README_CN.md b/README_CN.md index 5dad64af9b..8f4d4b05aa 100644 --- a/README_CN.md +++ b/README_CN.md @@ -20,7 +20,7 @@ Milvus 是一款开源的、针对海量特征向量的相似性搜索引擎。基于异构众核计算框架设计，成本更低，性能更好。在有限的计算资源下，十亿向量搜索仅毫秒响应。 -Milvus 提供稳定的 Python、Java 以及 C++ 的 API 接口。 +Milvus 提供稳定的 [Python](https://github.com/milvus-io/pymilvus)、[Java](https://github.com/milvus-io/milvus-sdk-java) 以及 [C++](https://github.com/milvus-io/milvus/tree/master/core/src/sdk) 的 API 接口。 通过 [版本发布说明](https://milvus.io/docs/zh-CN/release/v0.5.0/) 获取最新发行版本的 Milvus。 diff --git a/ci/jenkins/Jenkinsfile b/ci/jenkins/Jenkinsfile index 8d3953b112..bf7a70d11a 100644 --- a/ci/jenkins/Jenkinsfile +++ b/ci/jenkins/Jenkinsfile @@ -50,7 +50,7 @@ pipeline { } stages { - stage("Run Build") { + stage("Run GPU Version Build") { agent { kubernetes { label "${BINRARY_VERSION}-build" @@ -60,7 +60,7 @@ pipeline { } stages { - stage('Build') { + stage('GPU Version Build') { steps { container('milvus-build-env') { script { @@ -69,7 +69,7 @@ pipeline { } } } - stage('Code Coverage') { + stage('GPU Version Code Coverage') { steps { container('milvus-build-env') { script { @@ -78,7 +78,7 @@ pipeline { } } } - stage('Upload Package') { + stage('Upload GPU Version Package') { steps { container('milvus-build-env') { script { @@ -90,7 +90,7 @@ pipeline { } } - stage("Publish docker images") { + stage("Publish GPU Version docker images") { agent { kubernetes { label "${BINRARY_VERSION}-publish" @@ -100,7 +100,7 @@ pipeline { } stages { - stage('Publish') { + stage('Publish GPU Version') { steps { container('publish-images'){ script { @@ -112,7 +112,7 @@ pipeline { } } - stage("Deploy to Development") { + stage("Deploy GPU Version to Development") { agent { kubernetes { label "${BINRARY_VERSION}-dev-test" @@ -122,7 +122,7 @@ pipeline { } stages { - stage("Deploy to Dev") { + stage("Deploy GPU Version to Dev") { steps { container('milvus-test-env') { script { @@ -132,7 +132,7 @@ pipeline { } } - stage("Dev Test") { + stage("GPU Version Dev Test") { steps { container('milvus-test-env') { script { @@ -147,7 +147,7 @@ pipeline { } } - stage ("Cleanup Dev") { + stage ("Cleanup GPU Version Dev") { steps { container('milvus-test-env') { script { @@ -180,7 +180,7 @@ pipeline { } stages { - stage("Run Build") { + stage("Run CPU Version Build") { agent { kubernetes { label "${BINRARY_VERSION}-build" @@ -190,7 +190,7 @@ pipeline { } stages { - stage('Build') { + stage('Build CPU Version') { steps { container('milvus-build-env') { script { @@ -199,7 +199,7 @@ pipeline { } } } - stage('Code Coverage') { + stage('CPU Version Code Coverage') { steps { container('milvus-build-env') { script { @@ -208,7 +208,7 @@ pipeline { } } } - stage('Upload Package') { + stage('Upload CPU Version Package') { steps { container('milvus-build-env') { script { @@ -220,7 +220,7 @@ pipeline { } } - stage("Publish docker images") { + stage("Publish CPU Version docker images") { agent { kubernetes { label "${BINRARY_VERSION}-publish" @@ -230,7 +230,7 @@ pipeline { } stages { - stage('Publish') { + stage('Publish CPU Version') { steps { container('publish-images'){ script { @@ -242,7 +242,7 @@ pipeline { } } - stage("Deploy to Development") { + stage("Deploy CPU Version to Development") { agent { kubernetes { label "${BINRARY_VERSION}-dev-test" @@ -252,7 +252,7 @@ pipeline { } stages { - stage("Deploy to Dev") { + stage("Deploy CPU Version to Dev") { steps { container('milvus-test-env') { script { @@ -262,7 +262,7 @@ pipeline { } } - stage("Dev Test") { + stage("CPU Version Dev Test") { steps { container('milvus-test-env') { script { @@ -277,7 +277,7 @@ pipeline { } } - stage ("Cleanup Dev") { + stage ("Cleanup CPU Version Dev") { steps { container('milvus-test-env') { script { diff --git a/ci/scripts/build.sh b/ci/scripts/build.sh index 27962ccb54..eb5205ad26 100755 --- a/ci/scripts/build.sh +++ b/ci/scripts/build.sh @@ -16,6 +16,7 @@ BUILD_TYPE="Debug" BUILD_UNITTEST="OFF" INSTALL_PREFIX="/opt/milvus" FAISS_ROOT="" +CUSTOMIZATION="OFF" # default use origin faiss BUILD_COVERAGE="OFF" USE_JFROG_CACHE="OFF" RUN_CPPLINT="OFF" @@ -23,7 +24,7 @@ GPU_VERSION="OFF" WITH_MKL="OFF" CUDA_COMPILER=/usr/local/cuda/bin/nvcc -while getopts "o:t:b:f:gulcjmh" arg +while getopts "o:t:b:f:gxulcjmh" arg do case $arg in o) @@ -41,6 +42,9 @@ do g) GPU_VERSION="ON"; ;; + x) + CUSTOMIZATION="ON"; + ;; u) echo "Build and run unittest cases" ; BUILD_UNITTEST="ON"; @@ -66,6 +70,7 @@ parameter: -b: core code build directory -f: faiss root path -g: gpu version +-x: milvus customization (default: OFF) -u: building unit test options(default: OFF) -l: run cpplint, clang-format and clang-tidy(default: OFF) -c: code coverage(default: OFF) @@ -74,7 +79,7 @@ parameter: -h: help usage: -./build.sh -o \${INSTALL_PREFIX} -t \${BUILD_TYPE} -b \${CORE_BUILD_DIR} -f \${FAISS_ROOT} [-u] [-l] [-c] [-j] [-m] [-h] +./build.sh -o \${INSTALL_PREFIX} -t \${BUILD_TYPE} -b \${CORE_BUILD_DIR} -f \${FAISS_ROOT} [-g] [-x] [-u] [-l] [-c] [-j] [-m] [-h] " exit 0 ;; @@ -96,6 +101,7 @@ CMAKE_CMD="cmake \ -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \ -DCMAKE_CUDA_COMPILER=${CUDA_COMPILER} \ -DMILVUS_GPU_VERSION=${GPU_VERSION} \ +-DCUSTOMIZATION=${CUSTOMIZATION} \ -DBUILD_UNIT_TEST=${BUILD_UNITTEST} \ -DBUILD_COVERAGE=${BUILD_COVERAGE} \ -DUSE_JFROG_CACHE=${USE_JFROG_CACHE} \ diff --git a/core/.gitignore b/core/.gitignore index 74e41dba6b..8db8df41db 100644 --- a/core/.gitignore +++ b/core/.gitignore @@ -9,3 +9,5 @@ output.info output_new.info server.info *.pyc +src/grpc/python_gen.h +src/grpc/python/ diff --git a/core/conf/server_cpu_config.template b/core/conf/server_cpu_config.template index 1a384f19a1..41889f5cef 100644 --- a/core/conf/server_cpu_config.template +++ b/core/conf/server_cpu_config.template @@ -36,7 +36,7 @@ engine_config: gpu_search_threshold: 1000 # threshold beyond which the search computation is executed on GPUs only gpu_resource_config: - enable_gpu: false # whether to enable GPU resources + enable: false # whether to enable GPU resources cache_capacity: 4 # GB, size of GPU memory per card used for cache, must be a positive integer search_resources: # define the GPU devices used for search computation, must be in format gpux - gpu0 diff --git a/core/conf/server_gpu_config.template b/core/conf/server_gpu_config.template index 9790c93491..531c633da7 100644 --- a/core/conf/server_gpu_config.template +++ b/core/conf/server_gpu_config.template @@ -36,7 +36,7 @@ engine_config: gpu_search_threshold: 1000 # threshold beyond which the search computation is executed on GPUs only gpu_resource_config: - enable_gpu: true # whether to enable GPU resources + enable: true # whether to enable GPU resources cache_capacity: 4 # GB, size of GPU memory per card used for cache, must be a positive integer search_resources: # define the GPU devices used for search computation, must be in format gpux - gpu0 diff --git a/core/src/db/DBImpl.cpp b/core/src/db/DBImpl.cpp index 2559b3a46b..dd230ce0d1 100644 --- a/core/src/db/DBImpl.cpp +++ b/core/src/db/DBImpl.cpp @@ -84,12 +84,12 @@ DBImpl::Start() { return Status::OK(); } - ENGINE_LOG_TRACE << "DB service start"; + // ENGINE_LOG_TRACE << "DB service start"; shutting_down_.store(false, std::memory_order_release); // for distribute version, some nodes are read only if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) { - ENGINE_LOG_TRACE << "StartTimerTasks"; + // ENGINE_LOG_TRACE << "StartTimerTasks"; bg_timer_thread_ = std::thread(&DBImpl::BackgroundTimerTask, this); } @@ -114,7 +114,7 @@ DBImpl::Stop() { meta_ptr_->CleanUp(); } - ENGINE_LOG_TRACE << "DB service stop"; + // ENGINE_LOG_TRACE << "DB service stop"; return Status::OK(); } @@ -279,6 +279,11 @@ DBImpl::DropPartitionByTag(const std::string& table_id, const std::string& parti std::string partition_name; auto status = meta_ptr_->GetPartitionName(table_id, partition_tag, partition_name); + if (!status.ok()) { + ENGINE_LOG_ERROR << status.message(); + return status; + } + return DropPartition(partition_name); } @@ -553,7 +558,7 @@ DBImpl::StartMetricTask() { return; } - ENGINE_LOG_TRACE << "Start metric task"; + // ENGINE_LOG_TRACE << "Start metric task"; server::Metrics::GetInstance().KeepingAliveCounterIncrement(METRIC_ACTION_INTERVAL); int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage(); @@ -579,7 +584,7 @@ DBImpl::StartMetricTask() { server::Metrics::GetInstance().GPUTemperature(); server::Metrics::GetInstance().CPUTemperature(); - ENGINE_LOG_TRACE << "Metric task finished"; + // ENGINE_LOG_TRACE << "Metric task finished"; } Status @@ -751,7 +756,7 @@ DBImpl::BackgroundMergeFiles(const std::string& table_id) { void DBImpl::BackgroundCompaction(std::set table_ids) { - ENGINE_LOG_TRACE << "Background compaction thread start"; + // ENGINE_LOG_TRACE << " Background compaction thread start"; Status status; for (auto& table_id : table_ids) { @@ -774,7 +779,7 @@ DBImpl::BackgroundCompaction(std::set table_ids) { } meta_ptr_->CleanUpFilesWithTTL(ttl); - ENGINE_LOG_TRACE << "Background compaction thread exit"; + // ENGINE_LOG_TRACE << " Background compaction thread exit"; } void @@ -807,7 +812,7 @@ DBImpl::StartBuildIndexTask(bool force) { void DBImpl::BackgroundBuildIndex() { - ENGINE_LOG_TRACE << "Background build index thread start"; + // ENGINE_LOG_TRACE << "Background build index thread start"; std::unique_lock lock(build_index_mutex_); meta::TableFilesSchema to_index_files; @@ -830,7 +835,7 @@ DBImpl::BackgroundBuildIndex() { } } - ENGINE_LOG_TRACE << "Background build index thread exit"; + // ENGINE_LOG_TRACE << "Background build index thread exit"; } Status @@ -853,8 +858,12 @@ DBImpl::GetPartitionsByTags(const std::string& table_id, const std::vectorShowPartitions(table_id, partiton_array); for (auto& tag : partition_tags) { + // trim side-blank of tag, only compare valid characters + // for example: " ab cd " is treated as "ab cd" + std::string valid_tag = tag; + server::StringHelpFunctions::TrimStringBlank(valid_tag); for (auto& schema : partiton_array) { - if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, tag)) { + if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, valid_tag)) { partition_name_array.insert(schema.table_id_); } } diff --git a/core/src/db/Options.h b/core/src/db/Options.h index ebecb4de5a..91459a8d5f 100644 --- a/core/src/db/Options.h +++ b/core/src/db/Options.h @@ -33,7 +33,7 @@ static const char* ARCHIVE_CONF_DISK = "disk"; static const char* ARCHIVE_CONF_DAYS = "days"; struct ArchiveConf { - using CriteriaT = std::map; + using CriteriaT = std::map; explicit ArchiveConf(const std::string& type, const std::string& criterias = std::string()); diff --git a/core/src/db/engine/ExecutionEngine.h b/core/src/db/engine/ExecutionEngine.h index 86a014cf66..c8784e8a90 100644 --- a/core/src/db/engine/ExecutionEngine.h +++ b/core/src/db/engine/ExecutionEngine.h @@ -35,7 +35,9 @@ enum class EngineType { NSG_MIX, FAISS_IVFSQ8H, FAISS_PQ, - MAX_VALUE = FAISS_PQ, + SPTAG_KDT, + SPTAG_BKT, + MAX_VALUE = SPTAG_BKT, }; enum class MetricType { diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index 4f2f542254..ca307b90fc 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -124,6 +124,14 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) { #endif break; } + case EngineType::SPTAG_KDT: { + index = GetVecIndexFactory(IndexType::SPTAG_KDT_RNT_CPU); + break; + } + case EngineType::SPTAG_BKT: { + index = GetVecIndexFactory(IndexType::SPTAG_BKT_RNT_CPU); + break; + } default: { ENGINE_LOG_ERROR << "Unsupported index type"; return nullptr; @@ -146,7 +154,7 @@ ExecutionEngineImpl::HybridLoad() const { const std::string key = location_ + ".quantizer"; server::Config& config = server::Config::GetInstance(); - std::vector gpus; + std::vector gpus; Status s = config.GetGpuResourceConfigSearchResources(gpus); if (!s.ok()) { ENGINE_LOG_ERROR << s.message(); @@ -586,7 +594,7 @@ ExecutionEngineImpl::GpuCache(uint64_t gpu_id) { Status ExecutionEngineImpl::Init() { server::Config& config = server::Config::GetInstance(); - std::vector gpu_ids; + std::vector gpu_ids; Status s = config.GetGpuResourceConfigBuildIndexResources(gpu_ids); for (auto id : gpu_ids) { if (gpu_num_ == id) { diff --git a/core/src/db/engine/ExecutionEngineImpl.h b/core/src/db/engine/ExecutionEngineImpl.h index 7eb304426e..da0e7cfb64 100644 --- a/core/src/db/engine/ExecutionEngineImpl.h +++ b/core/src/db/engine/ExecutionEngineImpl.h @@ -122,8 +122,8 @@ class ExecutionEngineImpl : public ExecutionEngine { int64_t dim_; std::string location_; - int32_t nlist_ = 0; - int32_t gpu_num_ = 0; + int64_t nlist_ = 0; + int64_t gpu_num_ = 0; }; } // namespace engine diff --git a/core/src/db/meta/MySQLMetaImpl.cpp b/core/src/db/meta/MySQLMetaImpl.cpp index bf83447806..4406b87f7e 100644 --- a/core/src/db/meta/MySQLMetaImpl.cpp +++ b/core/src/db/meta/MySQLMetaImpl.cpp @@ -22,6 +22,7 @@ #include "metrics/Metrics.h" #include "utils/Exception.h" #include "utils/Log.h" +#include "utils/StringHelpFunctions.h" #include #include @@ -1162,17 +1163,23 @@ MySQLMetaImpl::CreatePartition(const std::string& table_id, const std::string& p // not allow create partition under partition if (!table_schema.owner_table_.empty()) { - return Status(DB_ERROR, "Nested partition is not allow"); + return Status(DB_ERROR, "Nested partition is not allowed"); + } + + // trim side-blank of tag, only compare valid characters + // for example: " ab cd " is treated as "ab cd" + std::string valid_tag = tag; + server::StringHelpFunctions::TrimStringBlank(valid_tag); + + // not allow duplicated partition + std::string exist_partition; + GetPartitionName(table_id, valid_tag, exist_partition); + if (!exist_partition.empty()) { + return Status(DB_ERROR, "Duplicate partition is not allowed"); } if (partition_name == "") { - // not allow duplicated partition - std::string exist_partition; - GetPartitionName(table_id, tag, exist_partition); - if (!exist_partition.empty()) { - return Status(DB_ERROR, "Duplicated partition is not allow"); - } - + // generate unique partition name NextTableId(table_schema.table_id_); } else { table_schema.table_id_ = partition_name; @@ -1182,9 +1189,14 @@ MySQLMetaImpl::CreatePartition(const std::string& table_id, const std::string& p table_schema.flag_ = 0; table_schema.created_on_ = utils::GetMicroSecTimeStamp(); table_schema.owner_table_ = table_id; - table_schema.partition_tag_ = tag; + table_schema.partition_tag_ = valid_tag; - return CreateTable(table_schema); + status = CreateTable(table_schema); + if (status.code() == DB_ALREADY_EXIST) { + return Status(DB_ALREADY_EXIST, "Partition already exists"); + } + + return status; } Status @@ -1231,6 +1243,12 @@ MySQLMetaImpl::GetPartitionName(const std::string& table_id, const std::string& try { server::MetricCollector metric; mysqlpp::StoreQueryResult res; + + // trim side-blank of tag, only compare valid characters + // for example: " ab cd " is treated as "ab cd" + std::string valid_tag = tag; + server::StringHelpFunctions::TrimStringBlank(valid_tag); + { mysqlpp::ScopedConnection connectionPtr(*mysql_connection_pool_, safe_grab_); @@ -1240,7 +1258,7 @@ MySQLMetaImpl::GetPartitionName(const std::string& table_id, const std::string& mysqlpp::Query allPartitionsQuery = connectionPtr->query(); allPartitionsQuery << "SELECT table_id FROM " << META_TABLES << " WHERE owner_table = " << mysqlpp::quote - << table_id << " AND partition_tag = " << mysqlpp::quote << tag << " AND state <> " + << table_id << " AND partition_tag = " << mysqlpp::quote << valid_tag << " AND state <> " << std::to_string(TableSchema::TO_DELETE) << ";"; ENGINE_LOG_DEBUG << "MySQLMetaImpl::AllTables: " << allPartitionsQuery.str(); @@ -1252,7 +1270,7 @@ MySQLMetaImpl::GetPartitionName(const std::string& table_id, const std::string& const mysqlpp::Row& resRow = res[0]; resRow["table_id"].to_string(partition_name); } else { - return Status(DB_NOT_FOUND, "Partition " + tag + " of table " + table_id + " not found"); + return Status(DB_NOT_FOUND, "Partition " + valid_tag + " of table " + table_id + " not found"); } } catch (std::exception& e) { return HandleException("GENERAL ERROR WHEN GET PARTITION NAME", e.what()); diff --git a/core/src/db/meta/SqliteMetaImpl.cpp b/core/src/db/meta/SqliteMetaImpl.cpp index 22e953fe9d..12128c074d 100644 --- a/core/src/db/meta/SqliteMetaImpl.cpp +++ b/core/src/db/meta/SqliteMetaImpl.cpp @@ -22,6 +22,7 @@ #include "metrics/Metrics.h" #include "utils/Exception.h" #include "utils/Log.h" +#include "utils/StringHelpFunctions.h" #include #include @@ -757,17 +758,23 @@ SqliteMetaImpl::CreatePartition(const std::string& table_id, const std::string& // not allow create partition under partition if(!table_schema.owner_table_.empty()) { - return Status(DB_ERROR, "Nested partition is not allow"); + return Status(DB_ERROR, "Nested partition is not allowed"); + } + + // trim side-blank of tag, only compare valid characters + // for example: " ab cd " is treated as "ab cd" + std::string valid_tag = tag; + server::StringHelpFunctions::TrimStringBlank(valid_tag); + + // not allow duplicated partition + std::string exist_partition; + GetPartitionName(table_id, valid_tag, exist_partition); + if(!exist_partition.empty()) { + return Status(DB_ERROR, "Duplicate partition is not allowed"); } if (partition_name == "") { - // not allow duplicated partition - std::string exist_partition; - GetPartitionName(table_id, tag, exist_partition); - if(!exist_partition.empty()) { - return Status(DB_ERROR, "Duplicated partition is not allow"); - } - + // generate unique partition name NextTableId(table_schema.table_id_); } else { table_schema.table_id_ = partition_name; @@ -777,9 +784,14 @@ SqliteMetaImpl::CreatePartition(const std::string& table_id, const std::string& table_schema.flag_ = 0; table_schema.created_on_ = utils::GetMicroSecTimeStamp(); table_schema.owner_table_ = table_id; - table_schema.partition_tag_ = tag; + table_schema.partition_tag_ = valid_tag; - return CreateTable(table_schema); + status = CreateTable(table_schema); + if (status.code() == DB_ALREADY_EXIST) { + return Status(DB_ALREADY_EXIST, "Partition already exists"); + } + + return status; } Status @@ -814,13 +826,18 @@ SqliteMetaImpl::GetPartitionName(const std::string& table_id, const std::string& try { server::MetricCollector metric; + // trim side-blank of tag, only compare valid characters + // for example: " ab cd " is treated as "ab cd" + std::string valid_tag = tag; + server::StringHelpFunctions::TrimStringBlank(valid_tag); + auto name = ConnectorPtr->select(columns(&TableSchema::table_id_), where(c(&TableSchema::owner_table_) == table_id - and c(&TableSchema::partition_tag_) == tag)); + and c(&TableSchema::partition_tag_) == valid_tag)); if (name.size() > 0) { partition_name = std::get<0>(name[0]); } else { - return Status(DB_NOT_FOUND, "Table " + table_id + "'s partition " + tag + " not found"); + return Status(DB_NOT_FOUND, "Table " + table_id + "'s partition " + valid_tag + " not found"); } } catch (std::exception &e) { return HandleException("Encounter exception when get partition name", e.what()); diff --git a/core/src/index/knowhere/CMakeLists.txt b/core/src/index/knowhere/CMakeLists.txt index 5f8e4d6970..285461bdef 100644 --- a/core/src/index/knowhere/CMakeLists.txt +++ b/core/src/index/knowhere/CMakeLists.txt @@ -30,10 +30,10 @@ set(external_srcs set(index_srcs knowhere/index/preprocessor/Normalize.cpp - knowhere/index/vector_index/IndexKDT.cpp + knowhere/index/vector_index/IndexSPTAG.cpp knowhere/index/vector_index/IndexIDMAP.cpp knowhere/index/vector_index/IndexIVF.cpp - knowhere/index/vector_index/helpers/KDTParameterMgr.cpp + knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp knowhere/index/vector_index/IndexNSG.cpp knowhere/index/vector_index/nsg/NSG.cpp knowhere/index/vector_index/nsg/NSGIO.cpp diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp deleted file mode 100644 index 1bd45075e3..0000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp +++ /dev/null @@ -1,180 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you 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 -#include -#include -#include -#include - -#undef mkdir - -#include "knowhere/index/vector_index/IndexKDT.h" -#include "knowhere/index/vector_index/helpers/Definitions.h" -//#include "knowhere/index/preprocessor/normalize.h" -#include "knowhere/adapter/SptagAdapter.h" -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h" - -namespace knowhere { - -BinarySet -CPUKDTRNG::Serialize() { - std::vector index_blobs; - std::vector index_len; - - // TODO(zirui): dev - // index_ptr_->SaveIndexToMemory(index_blobs, index_len); - BinarySet binary_set; - - // - // auto sample = std::make_shared(); - // sample.reset(static_cast(index_blobs[0])); - // auto tree = std::make_shared(); - // tree.reset(static_cast(index_blobs[1])); - // auto graph = std::make_shared(); - // graph.reset(static_cast(index_blobs[2])); - // auto metadata = std::make_shared(); - // metadata.reset(static_cast(index_blobs[3])); - // - // binary_set.Append("samples", sample, index_len[0]); - // binary_set.Append("tree", tree, index_len[1]); - // binary_set.Append("graph", graph, index_len[2]); - // binary_set.Append("metadata", metadata, index_len[3]); - return binary_set; -} - -void -CPUKDTRNG::Load(const BinarySet& binary_set) { - // TODO(zirui): dev - - // std::vector index_blobs; - // - // auto samples = binary_set.GetByName("samples"); - // index_blobs.push_back(samples->data.get()); - // - // auto tree = binary_set.GetByName("tree"); - // index_blobs.push_back(tree->data.get()); - // - // auto graph = binary_set.GetByName("graph"); - // index_blobs.push_back(graph->data.get()); - // - // auto metadata = binary_set.GetByName("metadata"); - // index_blobs.push_back(metadata->data.get()); - // - // index_ptr_->LoadIndexFromMemory(index_blobs); -} - -// PreprocessorPtr -// CPUKDTRNG::BuildPreprocessor(const DatasetPtr &dataset, const Config &config) { -// return std::make_shared(); -//} - -IndexModelPtr -CPUKDTRNG::Train(const DatasetPtr& origin, const Config& train_config) { - SetParameters(train_config); - DatasetPtr dataset = origin->Clone(); - - // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine - // && preprocessor_) { - // preprocessor_->Preprocess(dataset); - //} - - auto vectorset = ConvertToVectorSet(dataset); - auto metaset = ConvertToMetadataSet(dataset); - index_ptr_->BuildIndex(vectorset, metaset); - - // TODO: return IndexModelPtr - return nullptr; -} - -void -CPUKDTRNG::Add(const DatasetPtr& origin, const Config& add_config) { - SetParameters(add_config); - DatasetPtr dataset = origin->Clone(); - - // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine - // && preprocessor_) { - // preprocessor_->Preprocess(dataset); - //} - - auto vectorset = ConvertToVectorSet(dataset); - auto metaset = ConvertToMetadataSet(dataset); - index_ptr_->AddIndex(vectorset, metaset); -} - -void -CPUKDTRNG::SetParameters(const Config& config) { - for (auto& para : KDTParameterMgr::GetInstance().GetKDTParameters()) { - // auto value = config.get_with_default(para.first, para.second); - index_ptr_->SetParameter(para.first, para.second); - } -} - -DatasetPtr -CPUKDTRNG::Search(const DatasetPtr& dataset, const Config& config) { - SetParameters(config); - auto tensor = dataset->tensor()[0]; - auto p = (float*)tensor->raw_mutable_data(); - for (auto i = 0; i < 10; ++i) { - for (auto j = 0; j < 10; ++j) { - std::cout << p[i * 10 + j] << " "; - } - std::cout << std::endl; - } - std::vector query_results = ConvertToQueryResult(dataset, config); - -#pragma omp parallel for - for (auto i = 0; i < query_results.size(); ++i) { - auto target = (float*)query_results[i].GetTarget(); - std::cout << target[0] << ", " << target[1] << ", " << target[2] << std::endl; - index_ptr_->SearchIndex(query_results[i]); - } - - return ConvertToDataset(query_results); -} - -int64_t -CPUKDTRNG::Count() { - index_ptr_->GetNumSamples(); -} - -int64_t -CPUKDTRNG::Dimension() { - index_ptr_->GetFeatureDim(); -} - -VectorIndexPtr -CPUKDTRNG::Clone() { - KNOWHERE_THROW_MSG("not support"); -} - -void -CPUKDTRNG::Seal() { - // do nothing -} - -// TODO(linxj): -BinarySet -CPUKDTRNGIndexModel::Serialize() { -} - -void -CPUKDTRNGIndexModel::Load(const BinarySet& binary) { -} - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp new file mode 100644 index 0000000000..17a93fdcc7 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp @@ -0,0 +1,348 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 +#include +#include + +#include +#include +#include + +#undef mkdir + +#include "knowhere/adapter/SptagAdapter.h" +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/IndexSPTAG.h" +#include "knowhere/index/vector_index/helpers/Definitions.h" +#include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h" + +namespace knowhere { + +CPUSPTAGRNG::CPUSPTAGRNG(const std::string& IndexType) { + if (IndexType == "KDT") { + index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float); + index_ptr_->SetParameter("DistCalcMethod", "L2"); + index_type_ = SPTAG::IndexAlgoType::KDT; + } else { + index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::BKT, SPTAG::VectorValueType::Float); + index_ptr_->SetParameter("DistCalcMethod", "L2"); + index_type_ = SPTAG::IndexAlgoType::BKT; + } +} + +BinarySet +CPUSPTAGRNG::Serialize() { + std::string index_config; + std::vector index_blobs; + + std::shared_ptr> buffersize = index_ptr_->CalculateBufferSize(); + std::vector res(buffersize->size() + 1); + for (uint64_t i = 1; i < res.size(); i++) { + res[i] = new char[buffersize->at(i - 1)]; + auto ptr = &res[i][0]; + index_blobs.emplace_back(SPTAG::ByteArray((std::uint8_t*)ptr, buffersize->at(i - 1), false)); + } + + index_ptr_->SaveIndex(index_config, index_blobs); + + size_t length = index_config.length(); + char* cstr = new char[length]; + snprintf(cstr, length, "%s", index_config.c_str()); + + BinarySet binary_set; + auto sample = std::make_shared(); + sample.reset(static_cast(index_blobs[0].Data())); + auto tree = std::make_shared(); + tree.reset(static_cast(index_blobs[1].Data())); + auto graph = std::make_shared(); + graph.reset(static_cast(index_blobs[2].Data())); + auto deleteid = std::make_shared(); + deleteid.reset(static_cast(index_blobs[3].Data())); + auto metadata1 = std::make_shared(); + metadata1.reset(static_cast(index_blobs[4].Data())); + auto metadata2 = std::make_shared(); + metadata2.reset(static_cast(index_blobs[5].Data())); + auto config = std::make_shared(); + config.reset(static_cast((void*)cstr)); + + binary_set.Append("samples", sample, index_blobs[0].Length()); + binary_set.Append("tree", tree, index_blobs[1].Length()); + binary_set.Append("deleteid", deleteid, index_blobs[3].Length()); + binary_set.Append("metadata1", metadata1, index_blobs[4].Length()); + binary_set.Append("metadata2", metadata2, index_blobs[5].Length()); + binary_set.Append("config", config, length); + binary_set.Append("graph", graph, index_blobs[2].Length()); + + // MemoryIOWriter writer; + // size_t len = 0; + // for (int i = 0; i < 6; ++i) { + // len = index_blobs[i].Length(); + // assert(len != 0); + // writer(&len, sizeof(size_t), 1); + // writer(index_blobs[i].Data(), len, 1); + // len = 0; + // } + // writer(&length, sizeof(size_t), 1); + // writer(cstr, length, 1); + // auto data = std::make_shared(); + // data.reset(writer.data_); + // BinarySet binary_set; + // binary_set.Append("sptag", data, writer.total); + + // MemoryIOWriter writer; + // size_t len = 0; + // for (int i = 0; i < 6; ++i) { + // if (i == 2) continue; + // len = index_blobs[i].Length(); + // assert(len != 0); + // writer(&len, sizeof(size_t), 1); + // writer(index_blobs[i].Data(), len, 1); + // len = 0; + // } + // writer(&length, sizeof(size_t), 1); + // writer(cstr, length, 1); + // auto data = std::make_shared(); + // data.reset(writer.data_); + // BinarySet binary_set; + // binary_set.Append("sptag", data, writer.total); + // auto graph = std::make_shared(); + // graph.reset(static_cast(index_blobs[2].Data())); + // binary_set.Append("graph", graph, index_blobs[2].Length()); + + return binary_set; +} + +void +CPUSPTAGRNG::Load(const BinarySet& binary_set) { + std::string index_config; + std::vector index_blobs; + + auto samples = binary_set.GetByName("samples"); + index_blobs.push_back(SPTAG::ByteArray(samples->data.get(), samples->size, false)); + + auto tree = binary_set.GetByName("tree"); + index_blobs.push_back(SPTAG::ByteArray(tree->data.get(), tree->size, false)); + + auto graph = binary_set.GetByName("graph"); + index_blobs.push_back(SPTAG::ByteArray(graph->data.get(), graph->size, false)); + + auto deleteid = binary_set.GetByName("deleteid"); + index_blobs.push_back(SPTAG::ByteArray(deleteid->data.get(), deleteid->size, false)); + + auto metadata1 = binary_set.GetByName("metadata1"); + index_blobs.push_back(SPTAG::ByteArray(metadata1->data.get(), metadata1->size, false)); + + auto metadata2 = binary_set.GetByName("metadata2"); + index_blobs.push_back(SPTAG::ByteArray(metadata2->data.get(), metadata2->size, false)); + + auto config = binary_set.GetByName("config"); + index_config = reinterpret_cast(config->data.get()); + + // std::vector index_blobs; + // auto data = binary_set.GetByName("sptag"); + // MemoryIOReader reader; + // reader.total = data->size; + // reader.data_ = data->data.get(); + // size_t len = 0; + // for (int i = 0; i < 6; ++i) { + // reader(&len, sizeof(size_t), 1); + // assert(len != 0); + // auto binary = new uint8_t[len]; + // reader(binary, len, 1); + // index_blobs.emplace_back(SPTAG::ByteArray(binary, len, true)); + // len = 0; + // } + // reader(&len, sizeof(size_t), 1); + // assert(len != 0); + // auto config = new char[len]; + // reader(config, len, 1); + // std::string index_config = config; + // delete[] config; + + // std::vector index_blobs; + // auto data = binary_set.GetByName("sptag"); + // MemoryIOReader reader; + // reader.total = data->size; + // reader.data_ = data->data.get(); + // size_t len = 0; + // for (int i = 0; i < 6; ++i) { + // if (i == 2) { + // auto graph = binary_set.GetByName("graph"); + // index_blobs.emplace_back(SPTAG::ByteArray(graph->data.get(), graph->size, false)); + // continue; + // } + // reader(&len, sizeof(size_t), 1); + // assert(len != 0); + // auto binary = new uint8_t[len]; + // reader(binary, len, 1); + // index_blobs.emplace_back(SPTAG::ByteArray(binary, len, true)); + // len = 0; + // } + // reader(&len, sizeof(size_t), 1); + // assert(len != 0); + // auto config = new char[len]; + // reader(config, len, 1); + // std::string index_config = config; + // delete[] config; + index_ptr_->LoadIndex(index_config, index_blobs); +} + +// PreprocessorPtr +// CPUKDTRNG::BuildPreprocessor(const DatasetPtr &dataset, const Config &config) { +// return std::make_shared(); +//} + +IndexModelPtr +CPUSPTAGRNG::Train(const DatasetPtr& origin, const Config& train_config) { + SetParameters(train_config); + DatasetPtr dataset = origin->Clone(); + + // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine + // && preprocessor_) { + // preprocessor_->Preprocess(dataset); + //} + + auto vectorset = ConvertToVectorSet(dataset); + auto metaset = ConvertToMetadataSet(dataset); + index_ptr_->BuildIndex(vectorset, metaset); + + // TODO: return IndexModelPtr + return nullptr; +} + +void +CPUSPTAGRNG::Add(const DatasetPtr& origin, const Config& add_config) { + // SetParameters(add_config); + // DatasetPtr dataset = origin->Clone(); + // + // // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine + // // && preprocessor_) { + // // preprocessor_->Preprocess(dataset); + // //} + // + // auto vectorset = ConvertToVectorSet(dataset); + // auto metaset = ConvertToMetadataSet(dataset); + // index_ptr_->AddIndex(vectorset, metaset); +} + +void +CPUSPTAGRNG::SetParameters(const Config& config) { +#define Assign(param_name, str_name) \ + conf->param_name == INVALID_VALUE ? index_ptr_->SetParameter(str_name, std::to_string(build_cfg->param_name)) \ + : index_ptr_->SetParameter(str_name, std::to_string(conf->param_name)) + + if (index_type_ == SPTAG::IndexAlgoType::KDT) { + auto conf = std::dynamic_pointer_cast(config); + auto build_cfg = SPTAGParameterMgr::GetInstance().GetKDTParameters(); + + Assign(kdtnumber, "KDTNumber"); + Assign(numtopdimensionkdtsplit, "NumTopDimensionKDTSplit"); + Assign(samples, "Samples"); + Assign(tptnumber, "TPTNumber"); + Assign(tptleafsize, "TPTLeafSize"); + Assign(numtopdimensiontptsplit, "NumTopDimensionTPTSplit"); + Assign(neighborhoodsize, "NeighborhoodSize"); + Assign(graphneighborhoodscale, "GraphNeighborhoodScale"); + Assign(graphcefscale, "GraphCEFScale"); + Assign(refineiterations, "RefineIterations"); + Assign(cef, "CEF"); + Assign(maxcheckforrefinegraph, "MaxCheckForRefineGraph"); + Assign(numofthreads, "NumberOfThreads"); + Assign(maxcheck, "MaxCheck"); + Assign(thresholdofnumberofcontinuousnobetterpropagation, "ThresholdOfNumberOfContinuousNoBetterPropagation"); + Assign(numberofinitialdynamicpivots, "NumberOfInitialDynamicPivots"); + Assign(numberofotherdynamicpivots, "NumberOfOtherDynamicPivots"); + } else { + auto conf = std::dynamic_pointer_cast(config); + auto build_cfg = SPTAGParameterMgr::GetInstance().GetBKTParameters(); + + Assign(bktnumber, "BKTNumber"); + Assign(bktkmeansk, "BKTKMeansK"); + Assign(bktleafsize, "BKTLeafSize"); + Assign(samples, "Samples"); + Assign(tptnumber, "TPTNumber"); + Assign(tptleafsize, "TPTLeafSize"); + Assign(numtopdimensiontptsplit, "NumTopDimensionTPTSplit"); + Assign(neighborhoodsize, "NeighborhoodSize"); + Assign(graphneighborhoodscale, "GraphNeighborhoodScale"); + Assign(graphcefscale, "GraphCEFScale"); + Assign(refineiterations, "RefineIterations"); + Assign(cef, "CEF"); + Assign(maxcheckforrefinegraph, "MaxCheckForRefineGraph"); + Assign(numofthreads, "NumberOfThreads"); + Assign(maxcheck, "MaxCheck"); + Assign(thresholdofnumberofcontinuousnobetterpropagation, "ThresholdOfNumberOfContinuousNoBetterPropagation"); + Assign(numberofinitialdynamicpivots, "NumberOfInitialDynamicPivots"); + Assign(numberofotherdynamicpivots, "NumberOfOtherDynamicPivots"); + } +} + +DatasetPtr +CPUSPTAGRNG::Search(const DatasetPtr& dataset, const Config& config) { + SetParameters(config); + auto tensor = dataset->tensor()[0]; + auto p = (float*)tensor->raw_mutable_data(); + for (auto i = 0; i < 10; ++i) { + for (auto j = 0; j < 10; ++j) { + std::cout << p[i * 10 + j] << " "; + } + std::cout << std::endl; + } + std::vector query_results = ConvertToQueryResult(dataset, config); + +#pragma omp parallel for + for (auto i = 0; i < query_results.size(); ++i) { + auto target = (float*)query_results[i].GetTarget(); + std::cout << target[0] << ", " << target[1] << ", " << target[2] << std::endl; + index_ptr_->SearchIndex(query_results[i]); + } + + return ConvertToDataset(query_results); +} + +int64_t +CPUSPTAGRNG::Count() { + return index_ptr_->GetNumSamples(); +} + +int64_t +CPUSPTAGRNG::Dimension() { + return index_ptr_->GetFeatureDim(); +} + +VectorIndexPtr +CPUSPTAGRNG::Clone() { + KNOWHERE_THROW_MSG("not support"); +} + +void +CPUSPTAGRNG::Seal() { + return; // do nothing +} + +BinarySet +CPUSPTAGRNGIndexModel::Serialize() { + // KNOWHERE_THROW_MSG("not support"); // not support +} + +void +CPUSPTAGRNGIndexModel::Load(const BinarySet& binary) { + // KNOWHERE_THROW_MSG("not support"); // not support +} + +} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h similarity index 83% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h rename to core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h index f6d436995b..01380ce943 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h @@ -18,33 +18,37 @@ #pragma once #include + #include #include +#include + #include "VectorIndex.h" #include "knowhere/index/IndexModel.h" namespace knowhere { -class CPUKDTRNG : public VectorIndex { +class CPUSPTAGRNG : public VectorIndex { public: - CPUKDTRNG() { - index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float); - index_ptr_->SetParameter("DistCalcMethod", "L2"); - } + explicit CPUSPTAGRNG(const std::string& IndexType); public: BinarySet Serialize() override; + VectorIndexPtr Clone() override; + void Load(const BinarySet& index_array) override; public: // PreprocessorPtr // BuildPreprocessor(const DatasetPtr &dataset, const Config &config) override; + int64_t Count() override; + int64_t Dimension() override; @@ -56,6 +60,7 @@ class CPUKDTRNG : public VectorIndex { DatasetPtr Search(const DatasetPtr& dataset, const Config& config) override; + void Seal() override; @@ -66,11 +71,12 @@ class CPUKDTRNG : public VectorIndex { private: PreprocessorPtr preprocessor_; std::shared_ptr index_ptr_; + SPTAG::IndexAlgoType index_type_; }; -using CPUKDTRNGPtr = std::shared_ptr; +using CPUSPTAGRNGPtr = std::shared_ptr; -class CPUKDTRNGIndexModel : public IndexModel { +class CPUSPTAGRNGIndexModel : public IndexModel { public: BinarySet Serialize() override; @@ -82,6 +88,6 @@ class CPUKDTRNGIndexModel : public IndexModel { std::shared_ptr index_; }; -using CPUKDTRNGIndexModelPtr = std::shared_ptr; +using CPUSPTAGRNGIndexModelPtr = std::shared_ptr; } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h index b931790b04..e30088ecdf 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h @@ -42,6 +42,32 @@ constexpr int64_t DEFAULT_OUT_DEGREE = INVALID_VALUE; constexpr int64_t DEFAULT_CANDIDATE_SISE = INVALID_VALUE; constexpr int64_t DEFAULT_NNG_K = INVALID_VALUE; +// SPTAG Config +constexpr int64_t DEFAULT_SAMPLES = INVALID_VALUE; +constexpr int64_t DEFAULT_TPTNUMBER = INVALID_VALUE; +constexpr int64_t DEFAULT_TPTLEAFSIZE = INVALID_VALUE; +constexpr int64_t DEFAULT_NUMTOPDIMENSIONTPTSPLIT = INVALID_VALUE; +constexpr int64_t DEFAULT_NEIGHBORHOODSIZE = INVALID_VALUE; +constexpr int64_t DEFAULT_GRAPHNEIGHBORHOODSCALE = INVALID_VALUE; +constexpr int64_t DEFAULT_GRAPHCEFSCALE = INVALID_VALUE; +constexpr int64_t DEFAULT_REFINEITERATIONS = INVALID_VALUE; +constexpr int64_t DEFAULT_CEF = INVALID_VALUE; +constexpr int64_t DEFAULT_MAXCHECKFORREFINEGRAPH = INVALID_VALUE; +constexpr int64_t DEFAULT_NUMOFTHREADS = INVALID_VALUE; +constexpr int64_t DEFAULT_MAXCHECK = INVALID_VALUE; +constexpr int64_t DEFAULT_THRESHOLDOFNUMBEROFCONTINUOUSNOBETTERPROPAGATION = INVALID_VALUE; +constexpr int64_t DEFAULT_NUMBEROFINITIALDYNAMICPIVOTS = INVALID_VALUE; +constexpr int64_t DEFAULT_NUMBEROFOTHERDYNAMICPIVOTS = INVALID_VALUE; + +// KDT Config +constexpr int64_t DEFAULT_KDTNUMBER = INVALID_VALUE; +constexpr int64_t DEFAULT_NUMTOPDIMENSIONKDTSPLIT = INVALID_VALUE; + +// BKT Config +constexpr int64_t DEFAULT_BKTNUMBER = INVALID_VALUE; +constexpr int64_t DEFAULT_BKTKMEANSK = INVALID_VALUE; +constexpr int64_t DEFAULT_BKTLEAFSIZE = INVALID_VALUE; + struct IVFCfg : public Cfg { int64_t nlist = DEFAULT_NLIST; int64_t nprobe = DEFAULT_NPROBE; @@ -135,8 +161,57 @@ struct NSGCfg : public IVFCfg { }; using NSGConfig = std::shared_ptr; -struct KDTCfg : public Cfg { - int64_t tptnubmber = -1; +struct SPTAGCfg : public Cfg { + int64_t samples = DEFAULT_SAMPLES; + int64_t tptnumber = DEFAULT_TPTNUMBER; + int64_t tptleafsize = DEFAULT_TPTLEAFSIZE; + int64_t numtopdimensiontptsplit = DEFAULT_NUMTOPDIMENSIONTPTSPLIT; + int64_t neighborhoodsize = DEFAULT_NEIGHBORHOODSIZE; + int64_t graphneighborhoodscale = DEFAULT_GRAPHNEIGHBORHOODSCALE; + int64_t graphcefscale = DEFAULT_GRAPHCEFSCALE; + int64_t refineiterations = DEFAULT_REFINEITERATIONS; + int64_t cef = DEFAULT_CEF; + int64_t maxcheckforrefinegraph = DEFAULT_MAXCHECKFORREFINEGRAPH; + int64_t numofthreads = DEFAULT_NUMOFTHREADS; + int64_t maxcheck = DEFAULT_MAXCHECK; + int64_t thresholdofnumberofcontinuousnobetterpropagation = DEFAULT_THRESHOLDOFNUMBEROFCONTINUOUSNOBETTERPROPAGATION; + int64_t numberofinitialdynamicpivots = DEFAULT_NUMBEROFINITIALDYNAMICPIVOTS; + int64_t numberofotherdynamicpivots = DEFAULT_NUMBEROFOTHERDYNAMICPIVOTS; + + SPTAGCfg() = default; + + bool + CheckValid() override { + return true; + }; }; +using SPTAGConfig = std::shared_ptr; + +struct KDTCfg : public SPTAGCfg { + int64_t kdtnumber = DEFAULT_KDTNUMBER; + int64_t numtopdimensionkdtsplit = DEFAULT_NUMTOPDIMENSIONKDTSPLIT; + + KDTCfg() = default; + + bool + CheckValid() override { + return true; + }; +}; +using KDTConfig = std::shared_ptr; + +struct BKTCfg : public SPTAGCfg { + int64_t bktnumber = DEFAULT_BKTNUMBER; + int64_t bktkmeansk = DEFAULT_BKTKMEANSK; + int64_t bktleafsize = DEFAULT_BKTLEAFSIZE; + + BKTCfg() = default; + + bool + CheckValid() override { + return true; + }; +}; +using BKTConfig = std::shared_ptr; } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp deleted file mode 100644 index 19bf070dba..0000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp +++ /dev/null @@ -1,55 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you 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 - -#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h" - -namespace knowhere { - -const std::vector& -KDTParameterMgr::GetKDTParameters() { - return kdt_parameters_; -} - -KDTParameterMgr::KDTParameterMgr() { - kdt_parameters_ = std::vector{ - {"KDTNumber", "1"}, - {"NumTopDimensionKDTSplit", "5"}, - {"NumSamplesKDTSplitConsideration", "100"}, - - {"TPTNumber", "1"}, - {"TPTLeafSize", "2000"}, - {"NumTopDimensionTPTSplit", "5"}, - - {"NeighborhoodSize", "32"}, - {"GraphNeighborhoodScale", "2"}, - {"GraphCEFScale", "2"}, - {"RefineIterations", "0"}, - {"CEF", "1000"}, - {"MaxCheckForRefineGraph", "10000"}, - - {"NumberOfThreads", "1"}, - - {"MaxCheck", "8192"}, - {"ThresholdOfNumberOfContinuousNoBetterPropagation", "3"}, - {"NumberOfInitialDynamicPivots", "50"}, - {"NumberOfOtherDynamicPivots", "4"}, - }; -} - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp new file mode 100644 index 0000000000..836f204c77 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp @@ -0,0 +1,75 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 + +#include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h" + +namespace knowhere { + +const KDTConfig& +SPTAGParameterMgr::GetKDTParameters() { + return kdt_config_; +} + +const BKTConfig& +SPTAGParameterMgr::GetBKTParameters() { + return bkt_config_; +} + +SPTAGParameterMgr::SPTAGParameterMgr() { + kdt_config_ = std::make_shared(); + kdt_config_->kdtnumber = 1; + kdt_config_->numtopdimensionkdtsplit = 5; + kdt_config_->samples = 100; + kdt_config_->tptnumber = 1; + kdt_config_->tptleafsize = 2000; + kdt_config_->numtopdimensiontptsplit = 5; + kdt_config_->neighborhoodsize = 32; + kdt_config_->graphneighborhoodscale = 2; + kdt_config_->graphcefscale = 2; + kdt_config_->refineiterations = 0; + kdt_config_->cef = 1000; + kdt_config_->maxcheckforrefinegraph = 10000; + kdt_config_->numofthreads = 1; + kdt_config_->maxcheck = 8192; + kdt_config_->thresholdofnumberofcontinuousnobetterpropagation = 3; + kdt_config_->numberofinitialdynamicpivots = 50; + kdt_config_->numberofotherdynamicpivots = 4; + + bkt_config_ = std::make_shared(); + bkt_config_->bktnumber = 1; + bkt_config_->bktkmeansk = 32; + bkt_config_->bktleafsize = 8; + bkt_config_->samples = 100; + bkt_config_->tptnumber = 1; + bkt_config_->tptleafsize = 2000; + bkt_config_->numtopdimensiontptsplit = 5; + bkt_config_->neighborhoodsize = 32; + bkt_config_->graphneighborhoodscale = 2; + bkt_config_->graphcefscale = 2; + bkt_config_->refineiterations = 0; + bkt_config_->cef = 1000; + bkt_config_->maxcheckforrefinegraph = 10000; + bkt_config_->numofthreads = 1; + bkt_config_->maxcheck = 8192; + bkt_config_->thresholdofnumberofcontinuousnobetterpropagation = 3; + bkt_config_->numberofinitialdynamicpivots = 50; + bkt_config_->numberofotherdynamicpivots = 4; +} + +} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h similarity index 66% rename from core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h rename to core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h index fe90761e17..6a6f7c48d1 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h @@ -22,31 +22,40 @@ #include #include +#include +#include "IndexParameter.h" + namespace knowhere { -using KDTParameter = std::pair; +using KDTConfig = std::shared_ptr; +using BKTConfig = std::shared_ptr; -class KDTParameterMgr { +class SPTAGParameterMgr { public: - const std::vector& + const KDTConfig& GetKDTParameters(); + const BKTConfig& + GetBKTParameters(); + public: - static KDTParameterMgr& + static SPTAGParameterMgr& GetInstance() { - static KDTParameterMgr instance; + static SPTAGParameterMgr instance; return instance; } - KDTParameterMgr(const KDTParameterMgr&) = delete; - KDTParameterMgr& - operator=(const KDTParameterMgr&) = delete; + SPTAGParameterMgr(const SPTAGParameterMgr&) = delete; + + SPTAGParameterMgr& + operator=(const SPTAGParameterMgr&) = delete; private: - KDTParameterMgr(); + SPTAGParameterMgr(); private: - std::vector kdt_parameters_; + KDTConfig kdt_config_; + BKTConfig bkt_config_; }; } // namespace knowhere diff --git a/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h b/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h index d00ea45365..0208f6d983 100644 --- a/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h +++ b/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h @@ -195,7 +195,7 @@ namespace SPTAG C = *((DimensionType*)pDataPointsMemFile); pDataPointsMemFile += sizeof(DimensionType); - Initialize(R, C, (T*)pDataPointsMemFile); + Initialize(R, C, (T*)pDataPointsMemFile, false); std::cout << "Load " << name << " (" << R << ", " << C << ") Finish!" << std::endl; return true; } diff --git a/core/src/index/unittest/CMakeLists.txt b/core/src/index/unittest/CMakeLists.txt index d3ba84e557..93ae63a9ec 100644 --- a/core/src/index/unittest/CMakeLists.txt +++ b/core/src/index/unittest/CMakeLists.txt @@ -82,17 +82,17 @@ if (NOT TARGET test_idmap) endif () target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs}) -# -set(kdt_srcs +# +set(sptag_srcs ${INDEX_SOURCE_DIR}/knowhere/knowhere/adapter/SptagAdapter.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/preprocessor/Normalize.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexKDT.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp ) -if (NOT TARGET test_kdt) - add_executable(test_kdt test_kdt.cpp ${kdt_srcs} ${util_srcs}) +if (NOT TARGET test_sptag) + add_executable(test_sptag test_sptag.cpp ${sptag_srcs} ${util_srcs}) endif () -target_link_libraries(test_kdt +target_link_libraries(test_sptag SPTAGLibStatic ${depend_libs} ${unittest_libs} ${basic_libs}) @@ -106,7 +106,7 @@ endif () install(TARGETS test_ivf DESTINATION unittest) install(TARGETS test_idmap DESTINATION unittest) -install(TARGETS test_kdt DESTINATION unittest) +install(TARGETS test_sptag DESTINATION unittest) if (KNOWHERE_GPU_VERSION) install(TARGETS test_gpuresource DESTINATION unittest) install(TARGETS test_customized_index DESTINATION unittest) diff --git a/core/src/index/unittest/test_kdt.cpp b/core/src/index/unittest/test_kdt.cpp deleted file mode 100644 index bbc7dcf94c..0000000000 --- a/core/src/index/unittest/test_kdt.cpp +++ /dev/null @@ -1,144 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you 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 - -#include -#include - -#include "knowhere/adapter/SptagAdapter.h" -#include "knowhere/adapter/Structure.h" -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexKDT.h" -#include "knowhere/index/vector_index/helpers/Definitions.h" - -#include "unittest/utils.h" - -using ::testing::Combine; -using ::testing::TestWithParam; -using ::testing::Values; - -class KDTTest : public DataGen, public ::testing::Test { - protected: - void - SetUp() override { - Generate(96, 1000, 10); - index_ = std::make_shared(); - - auto tempconf = std::make_shared(); - tempconf->tptnubmber = 1; - tempconf->k = 10; - conf = tempconf; - - Init_with_default(); - } - - protected: - knowhere::Config conf; - std::shared_ptr index_ = nullptr; -}; - -// TODO(lxj): add test about count() and dimension() -TEST_F(KDTTest, kdt_basic) { - assert(!xb.empty()); - - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, k); - - { - // auto ids = result->array()[0]; - // auto dists = result->array()[1]; - auto ids = result->ids(); - auto dists = result->dist(); - - std::stringstream ss_id; - std::stringstream ss_dist; - for (auto i = 0; i < nq; i++) { - for (auto j = 0; j < k; ++j) { - ss_id << *((int64_t*)(ids) + i * k + j) << " "; - ss_dist << *((float*)(dists) + i * k + j) << " "; - // ss_id << *ids->data()->GetValues(1, i * k + j) << " "; - // ss_dist << *dists->data()->GetValues(1, i * k + j) << " "; - } - ss_id << std::endl; - ss_dist << std::endl; - } - std::cout << "id\n" << ss_id.str() << std::endl; - std::cout << "dist\n" << ss_dist.str() << std::endl; - } -} - -// TODO(zirui): enable test -// TEST_F(KDTTest, kdt_serialize) { -// assert(!xb.empty()); -// -// auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); -// index_->set_preprocessor(preprocessor); -// -// auto model = index_->Train(base_dataset, conf); -// // index_->Add(base_dataset, conf); -// auto binaryset = index_->Serialize(); -// auto new_index = std::make_shared(); -// new_index->Load(binaryset); -// auto result = new_index->Search(query_dataset, conf); -// AssertAnns(result, nq, k); -// PrintResult(result, nq, k); -// ASSERT_EQ(new_index->Count(), nb); -// ASSERT_EQ(new_index->Dimension(), dim); -// ASSERT_THROW({ new_index->Clone(); }, knowhere::KnowhereException); -// ASSERT_NO_THROW({ new_index->Seal(); }); -// -// { -// int fileno = 0; -// const std::string& base_name = "/tmp/kdt_serialize_test_bin_"; -// std::vector filename_list; -// std::vector> meta_list; -// for (auto& iter : binaryset.binary_map_) { -// const std::string& filename = base_name + std::to_string(fileno); -// FileIOWriter writer(filename); -// writer(iter.second->data.get(), iter.second->size); -// -// meta_list.emplace_back(std::make_pair(iter.first, iter.second->size)); -// filename_list.push_back(filename); -// ++fileno; -// } -// -// knowhere::BinarySet load_data_list; -// for (int i = 0; i < filename_list.size() && i < meta_list.size(); ++i) { -// auto bin_size = meta_list[i].second; -// FileIOReader reader(filename_list[i]); -// -// auto load_data = new uint8_t[bin_size]; -// reader(load_data, bin_size); -// auto data = std::make_shared(); -// data.reset(load_data); -// load_data_list.Append(meta_list[i].first, data, bin_size); -// } -// -// auto new_index = std::make_shared(); -// new_index->Load(load_data_list); -// auto result = new_index->Search(query_dataset, conf); -// AssertAnns(result, nq, k); -// PrintResult(result, nq, k); -// } -//} diff --git a/core/src/index/unittest/test_sptag.cpp b/core/src/index/unittest/test_sptag.cpp new file mode 100644 index 0000000000..d472dab0d7 --- /dev/null +++ b/core/src/index/unittest/test_sptag.cpp @@ -0,0 +1,154 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 + +#include +#include + +#include "knowhere/adapter/SptagAdapter.h" +#include "knowhere/adapter/Structure.h" +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/IndexSPTAG.h" +#include "knowhere/index/vector_index/helpers/Definitions.h" + +#include "unittest/utils.h" + +using ::testing::Combine; +using ::testing::TestWithParam; +using ::testing::Values; + +class SPTAGTest : public DataGen, public TestWithParam { + protected: + void + SetUp() override { + IndexType = GetParam(); + Generate(128, 100, 5); + index_ = std::make_shared(IndexType); + if (IndexType == "KDT") { + auto tempconf = std::make_shared(); + tempconf->tptnumber = 1; + tempconf->k = 10; + conf = tempconf; + } else { + auto tempconf = std::make_shared(); + tempconf->tptnumber = 1; + tempconf->k = 10; + conf = tempconf; + } + + Init_with_default(); + } + + protected: + knowhere::Config conf; + std::shared_ptr index_ = nullptr; + std::string IndexType; +}; + +INSTANTIATE_TEST_CASE_P(SPTAGParameters, SPTAGTest, Values("KDT", "BKT")); + +// TODO(lxj): add test about count() and dimension() +TEST_P(SPTAGTest, sptag_basic) { + assert(!xb.empty()); + + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + + auto model = index_->Train(base_dataset, conf); + index_->set_index_model(model); + index_->Add(base_dataset, conf); + auto result = index_->Search(query_dataset, conf); + AssertAnns(result, nq, k); + + { + // auto ids = result->array()[0]; + // auto dists = result->array()[1]; + auto ids = result->ids(); + auto dists = result->dist(); + + std::stringstream ss_id; + std::stringstream ss_dist; + for (auto i = 0; i < nq; i++) { + for (auto j = 0; j < k; ++j) { + // ss_id << *ids->data()->GetValues(1, i * k + j) << " "; + // ss_dist << *dists->data()->GetValues(1, i * k + j) << " "; + ss_id << *((int64_t*)(ids) + i * k + j) << " "; + ss_dist << *((float*)(dists) + i * k + j) << " "; + } + ss_id << std::endl; + ss_dist << std::endl; + } + std::cout << "id\n" << ss_id.str() << std::endl; + std::cout << "dist\n" << ss_dist.str() << std::endl; + } +} + +TEST_P(SPTAGTest, sptag_serialize) { + assert(!xb.empty()); + + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + + auto model = index_->Train(base_dataset, conf); + + index_->Add(base_dataset, conf); + auto binaryset = index_->Serialize(); + auto new_index = std::make_shared(IndexType); + new_index->Load(binaryset); + auto result = new_index->Search(query_dataset, conf); + AssertAnns(result, nq, k); + PrintResult(result, nq, k); + ASSERT_EQ(new_index->Count(), nb); + ASSERT_EQ(new_index->Dimension(), dim); + // ASSERT_THROW({ new_index->Clone(); }, knowhere::KnowhereException); + // ASSERT_NO_THROW({ new_index->Seal(); }); + + { + int fileno = 0; + const std::string& base_name = "/tmp/sptag_serialize_test_bin_"; + std::vector filename_list; + std::vector> meta_list; + for (auto& iter : binaryset.binary_map_) { + const std::string& filename = base_name + std::to_string(fileno); + FileIOWriter writer(filename); + writer(iter.second->data.get(), iter.second->size); + + meta_list.emplace_back(std::make_pair(iter.first, iter.second->size)); + filename_list.push_back(filename); + ++fileno; + } + + knowhere::BinarySet load_data_list; + for (int i = 0; i < filename_list.size() && i < meta_list.size(); ++i) { + auto bin_size = meta_list[i].second; + FileIOReader reader(filename_list[i]); + + auto load_data = new uint8_t[bin_size]; + reader(load_data, bin_size); + auto data = std::make_shared(); + data.reset(load_data); + load_data_list.Append(meta_list[i].first, data, bin_size); + } + + auto new_index = std::make_shared(IndexType); + new_index->Load(load_data_list); + auto result = new_index->Search(query_dataset, conf); + AssertAnns(result, nq, k); + PrintResult(result, nq, k); + } +} diff --git a/core/src/index/unittest/utils.cpp b/core/src/index/unittest/utils.cpp index 2556b60fad..11dad4a8b9 100644 --- a/core/src/index/unittest/utils.cpp +++ b/core/src/index/unittest/utils.cpp @@ -160,15 +160,17 @@ AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { void PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - auto dists = result->array()[1]; + auto ids = result->ids(); + auto dists = result->dist(); std::stringstream ss_id; std::stringstream ss_dist; - for (auto i = 0; i < 10; i++) { + for (auto i = 0; i < nq; i++) { for (auto j = 0; j < k; ++j) { - ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; + // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; + // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; + ss_id << *((int64_t*)(ids) + i * k + j) << " "; + ss_dist << *((float*)(dists) + i * k + j) << " "; } ss_id << std::endl; ss_dist << std::endl; diff --git a/core/src/scheduler/SchedInst.cpp b/core/src/scheduler/SchedInst.cpp index 867052ef5d..69d293f986 100644 --- a/core/src/scheduler/SchedInst.cpp +++ b/core/src/scheduler/SchedInst.cpp @@ -55,9 +55,9 @@ load_simple_config() { // get resources #ifdef MILVUS_GPU_VERSION server::Config& config = server::Config::GetInstance(); - std::vector gpu_ids; + std::vector gpu_ids; config.GetGpuResourceConfigSearchResources(gpu_ids); - std::vector build_gpu_ids; + std::vector build_gpu_ids; config.GetGpuResourceConfigBuildIndexResources(build_gpu_ids); auto pcie = Connection("pcie", 12000); diff --git a/core/src/scheduler/optimizer/BuildIndexPass.cpp b/core/src/scheduler/optimizer/BuildIndexPass.cpp index cd1b7550d8..d535b9675f 100644 --- a/core/src/scheduler/optimizer/BuildIndexPass.cpp +++ b/core/src/scheduler/optimizer/BuildIndexPass.cpp @@ -26,7 +26,7 @@ namespace scheduler { void BuildIndexPass::Init() { server::Config& config = server::Config::GetInstance(); - std::vector build_resources; + std::vector build_resources; Status s = config.GetGpuResourceConfigBuildIndexResources(build_resources); if (!s.ok()) { throw; diff --git a/core/src/scheduler/optimizer/FaissFlatPass.h b/core/src/scheduler/optimizer/FaissFlatPass.h index 60adf62a07..f219bebdf3 100644 --- a/core/src/scheduler/optimizer/FaissFlatPass.h +++ b/core/src/scheduler/optimizer/FaissFlatPass.h @@ -45,9 +45,9 @@ class FaissFlatPass : public Pass { Run(const TaskPtr& task) override; private: - int32_t threshold_ = std::numeric_limits::max(); + int64_t threshold_ = std::numeric_limits::max(); int64_t count_ = 0; - std::vector gpus; + std::vector gpus; }; using FaissFlatPassPtr = std::shared_ptr; diff --git a/core/src/scheduler/optimizer/FaissIVFFlatPass.h b/core/src/scheduler/optimizer/FaissIVFFlatPass.h index 9264db32a8..2d15539014 100644 --- a/core/src/scheduler/optimizer/FaissIVFFlatPass.h +++ b/core/src/scheduler/optimizer/FaissIVFFlatPass.h @@ -45,9 +45,9 @@ class FaissIVFFlatPass : public Pass { Run(const TaskPtr& task) override; private: - int32_t threshold_ = std::numeric_limits::max(); + int64_t threshold_ = std::numeric_limits::max(); int64_t count_ = 0; - std::vector gpus; + std::vector gpus; }; using FaissIVFFlatPassPtr = std::shared_ptr; diff --git a/core/src/scheduler/optimizer/FaissIVFSQ8HPass.cpp b/core/src/scheduler/optimizer/FaissIVFSQ8HPass.cpp index b63a0f7490..a99e861e03 100644 --- a/core/src/scheduler/optimizer/FaissIVFSQ8HPass.cpp +++ b/core/src/scheduler/optimizer/FaissIVFSQ8HPass.cpp @@ -32,7 +32,7 @@ FaissIVFSQ8HPass::Init() { server::Config& config = server::Config::GetInstance(); Status s = config.GetEngineConfigGpuSearchThreshold(threshold_); if (!s.ok()) { - threshold_ = std::numeric_limits::max(); + threshold_ = std::numeric_limits::max(); } s = config.GetGpuResourceConfigSearchResources(gpus); } diff --git a/core/src/scheduler/optimizer/FaissIVFSQ8HPass.h b/core/src/scheduler/optimizer/FaissIVFSQ8HPass.h index fa830ff08a..0d2892809f 100644 --- a/core/src/scheduler/optimizer/FaissIVFSQ8HPass.h +++ b/core/src/scheduler/optimizer/FaissIVFSQ8HPass.h @@ -45,9 +45,9 @@ class FaissIVFSQ8HPass : public Pass { Run(const TaskPtr& task) override; private: - int32_t threshold_ = std::numeric_limits::max(); + int64_t threshold_ = std::numeric_limits::max(); int64_t count_ = 0; - std::vector gpus; + std::vector gpus; }; using FaissIVFSQ8HPassPtr = std::shared_ptr; diff --git a/core/src/scheduler/optimizer/FaissIVFSQ8Pass.h b/core/src/scheduler/optimizer/FaissIVFSQ8Pass.h index c5853283f7..e92ea2fe4e 100644 --- a/core/src/scheduler/optimizer/FaissIVFSQ8Pass.h +++ b/core/src/scheduler/optimizer/FaissIVFSQ8Pass.h @@ -45,9 +45,9 @@ class FaissIVFSQ8Pass : public Pass { Run(const TaskPtr& task) override; private: - int32_t threshold_ = std::numeric_limits::max(); + int64_t threshold_ = std::numeric_limits::max(); int64_t count_ = 0; - std::vector gpus; + std::vector gpus; }; using FaissIVFSQ8PassPtr = std::shared_ptr; diff --git a/core/src/scheduler/task/BuildIndexTask.cpp b/core/src/scheduler/task/BuildIndexTask.cpp index d8602c141e..f561fa947d 100644 --- a/core/src/scheduler/task/BuildIndexTask.cpp +++ b/core/src/scheduler/task/BuildIndexTask.cpp @@ -146,8 +146,7 @@ XBuildIndexTask::Execute() { status = meta_ptr->UpdateTableFile(table_file); ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete"; - std::cout << "ERROR: failed to build index, index file is too large or gpu memory is not enough" - << std::endl; + ENGINE_LOG_ERROR << "Failed to build index, index file is too large or gpu memory is not enough"; build_index_job->BuildIndexDone(to_index_id_); build_index_job->GetStatus() = Status(DB_ERROR, msg); @@ -179,8 +178,8 @@ XBuildIndexTask::Execute() { status = meta_ptr->UpdateTableFile(table_file); ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete"; - std::cout << "ERROR: failed to persist index file: " << table_file.location_ - << ", possible out of disk space" << std::endl; + ENGINE_LOG_ERROR << "Failed to persist index file: " << table_file.location_ + << ", possible out of disk space"; build_index_job->BuildIndexDone(to_index_id_); build_index_job->GetStatus() = Status(DB_ERROR, msg); diff --git a/core/src/server/Config.cpp b/core/src/server/Config.cpp index e3a48b6ae4..f3efcff0cc 100644 --- a/core/src/server/Config.cpp +++ b/core/src/server/Config.cpp @@ -113,19 +113,19 @@ Config::ValidateConfig() { return s; } - int32_t db_archive_disk_threshold; + int64_t db_archive_disk_threshold; s = GetDBConfigArchiveDiskThreshold(db_archive_disk_threshold); if (!s.ok()) { return s; } - int32_t db_archive_days_threshold; + int64_t db_archive_days_threshold; s = GetDBConfigArchiveDaysThreshold(db_archive_days_threshold); if (!s.ok()) { return s; } - int32_t db_insert_buffer_size; + int64_t db_insert_buffer_size; s = GetDBConfigInsertBufferSize(db_insert_buffer_size); if (!s.ok()) { return s; @@ -170,19 +170,19 @@ Config::ValidateConfig() { } /* engine config */ - int32_t engine_use_blas_threshold; + int64_t engine_use_blas_threshold; s = GetEngineConfigUseBlasThreshold(engine_use_blas_threshold); if (!s.ok()) { return s; } - int32_t engine_omp_thread_num; + int64_t engine_omp_thread_num; s = GetEngineConfigOmpThreadNum(engine_omp_thread_num); if (!s.ok()) { return s; } - int32_t engine_gpu_search_threshold; + int64_t engine_gpu_search_threshold; s = GetEngineConfigGpuSearchThreshold(engine_gpu_search_threshold); if (!s.ok()) { return s; @@ -190,8 +190,8 @@ Config::ValidateConfig() { /* gpu resource config */ #ifdef MILVUS_GPU_VERSION - bool resource_enable_gpu; - s = GetGpuResourceConfigEnableGpu(resource_enable_gpu); + bool gpu_resource_enable; + s = GetGpuResourceConfigEnable(gpu_resource_enable); if (!s.ok()) { return s; } @@ -208,13 +208,13 @@ Config::ValidateConfig() { return s; } - std::vector search_resources; + std::vector search_resources; s = GetGpuResourceConfigSearchResources(search_resources); if (!s.ok()) { return s; } - std::vector index_build_resources; + std::vector index_build_resources; s = GetGpuResourceConfigBuildIndexResources(index_build_resources); if (!s.ok()) { return s; @@ -330,7 +330,7 @@ Config::ResetDefaultConfig() { /* gpu resource config */ #ifdef MILVUS_GPU_VERSION - s = SetGpuResourceConfigEnableGpu(CONFIG_GPU_RESOURCE_ENABLE_GPU_DEFAULT); + s = SetGpuResourceConfigEnable(CONFIG_GPU_RESOURCE_ENABLE_DEFAULT); if (!s.ok()) { return s; } @@ -485,7 +485,7 @@ Config::CheckDBConfigInsertBufferSize(const std::string& value) { ". Possible reason: db_config.insert_buffer_size is not a positive integer."; return Status(SERVER_INVALID_ARGUMENT, msg); } else { - int64_t buffer_size = std::stoi(value) * GB; + int64_t buffer_size = std::stoll(value) * GB; if (buffer_size <= 0) { std::string msg = "Invalid insert buffer size: " + value + ". Possible reason: db_config.insert_buffer_size is not a positive integer."; @@ -540,7 +540,7 @@ Config::CheckCacheConfigCpuCacheCapacity(const std::string& value) { ". Possible reason: cache_config.cpu_cache_capacity is not a positive integer."; return Status(SERVER_INVALID_ARGUMENT, msg); } else { - int64_t cpu_cache_capacity = std::stoi(value) * GB; + int64_t cpu_cache_capacity = std::stoll(value) * GB; if (cpu_cache_capacity <= 0) { std::string msg = "Invalid cpu cache capacity: " + value + ". Possible reason: cache_config.cpu_cache_capacity is not a positive integer."; @@ -557,7 +557,7 @@ Config::CheckCacheConfigCpuCacheCapacity(const std::string& value) { std::cerr << "WARNING: cpu cache capacity value is too big" << std::endl; } - int32_t buffer_value; + int64_t buffer_value; Status s = GetDBConfigInsertBufferSize(buffer_value); if (!s.ok()) { return s; @@ -619,10 +619,10 @@ Config::CheckEngineConfigOmpThreadNum(const std::string& value) { return Status(SERVER_INVALID_ARGUMENT, msg); } - int32_t omp_thread = std::stoi(value); - uint32_t sys_thread_cnt = 8; + int64_t omp_thread = std::stoll(value); + int64_t sys_thread_cnt = 8; CommonUtil::GetSystemAvailableThreads(sys_thread_cnt); - if (omp_thread > static_cast(sys_thread_cnt)) { + if (omp_thread > sys_thread_cnt) { std::string msg = "Invalid omp thread num: " + value + ". Possible reason: engine_config.omp_thread_num exceeds system cpu cores."; return Status(SERVER_INVALID_ARGUMENT, msg); @@ -641,10 +641,10 @@ Config::CheckEngineConfigGpuSearchThreshold(const std::string& value) { } Status -Config::CheckGpuResourceConfigEnableGpu(const std::string& value) { +Config::CheckGpuResourceConfigEnable(const std::string& value) { if (!ValidationUtil::ValidateStringIsBool(value).ok()) { - std::string msg = "Invalid gpu resource config: " + value + - ". Possible reason: gpu_resource_config.enable_gpu is not a boolean."; + std::string msg = + "Invalid gpu resource config: " + value + ". Possible reason: gpu_resource_config.enable is not a boolean."; return Status(SERVER_INVALID_ARGUMENT, msg); } return Status::OK(); @@ -657,14 +657,14 @@ Config::CheckGpuResourceConfigCacheCapacity(const std::string& value) { ". Possible reason: gpu_resource_config.cache_capacity is not a positive integer."; return Status(SERVER_INVALID_ARGUMENT, msg); } else { - uint64_t gpu_cache_capacity = std::stoi(value) * GB; - std::vector gpu_ids; + int64_t gpu_cache_capacity = std::stoll(value) * GB; + std::vector gpu_ids; Status s = GetGpuResourceConfigBuildIndexResources(gpu_ids); if (!s.ok()) { return s; } - for (int32_t gpu_id : gpu_ids) { + for (int64_t gpu_id : gpu_ids) { size_t gpu_memory; if (!ValidationUtil::GetGpuMemory(gpu_id, gpu_memory).ok()) { std::string msg = "Fail to get GPU memory for GPU device: " + std::to_string(gpu_id); @@ -855,37 +855,37 @@ Config::GetDBConfigBackendUrl(std::string& value) { } Status -Config::GetDBConfigArchiveDiskThreshold(int32_t& value) { +Config::GetDBConfigArchiveDiskThreshold(int64_t& value) { std::string str = GetConfigStr(CONFIG_DB, CONFIG_DB_ARCHIVE_DISK_THRESHOLD, CONFIG_DB_ARCHIVE_DISK_THRESHOLD_DEFAULT); Status s = CheckDBConfigArchiveDiskThreshold(str); if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } Status -Config::GetDBConfigArchiveDaysThreshold(int32_t& value) { +Config::GetDBConfigArchiveDaysThreshold(int64_t& value) { std::string str = GetConfigStr(CONFIG_DB, CONFIG_DB_ARCHIVE_DAYS_THRESHOLD, CONFIG_DB_ARCHIVE_DAYS_THRESHOLD_DEFAULT); Status s = CheckDBConfigArchiveDaysThreshold(str); if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } Status -Config::GetDBConfigInsertBufferSize(int32_t& value) { +Config::GetDBConfigInsertBufferSize(int64_t& value) { std::string str = GetConfigStr(CONFIG_DB, CONFIG_DB_INSERT_BUFFER_SIZE, CONFIG_DB_INSERT_BUFFER_SIZE_DEFAULT); Status s = CheckDBConfigInsertBufferSize(str); if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } @@ -927,7 +927,7 @@ Config::GetCacheConfigCpuCacheCapacity(int64_t& value) { if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } @@ -957,45 +957,44 @@ Config::GetCacheConfigCacheInsertData(bool& value) { } Status -Config::GetEngineConfigUseBlasThreshold(int32_t& value) { +Config::GetEngineConfigUseBlasThreshold(int64_t& value) { std::string str = GetConfigStr(CONFIG_ENGINE, CONFIG_ENGINE_USE_BLAS_THRESHOLD, CONFIG_ENGINE_USE_BLAS_THRESHOLD_DEFAULT); Status s = CheckEngineConfigUseBlasThreshold(str); if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } Status -Config::GetEngineConfigOmpThreadNum(int32_t& value) { +Config::GetEngineConfigOmpThreadNum(int64_t& value) { std::string str = GetConfigStr(CONFIG_ENGINE, CONFIG_ENGINE_OMP_THREAD_NUM, CONFIG_ENGINE_OMP_THREAD_NUM_DEFAULT); Status s = CheckEngineConfigOmpThreadNum(str); if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } Status -Config::GetEngineConfigGpuSearchThreshold(int32_t& value) { +Config::GetEngineConfigGpuSearchThreshold(int64_t& value) { std::string str = GetConfigStr(CONFIG_ENGINE, CONFIG_ENGINE_GPU_SEARCH_THRESHOLD, CONFIG_ENGINE_GPU_SEARCH_THRESHOLD_DEFAULT); Status s = CheckEngineConfigGpuSearchThreshold(str); if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } Status -Config::GetGpuResourceConfigEnableGpu(bool& value) { - std::string str = - GetConfigStr(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_ENABLE_GPU, CONFIG_GPU_RESOURCE_ENABLE_GPU_DEFAULT); - Status s = CheckGpuResourceConfigEnableGpu(str); +Config::GetGpuResourceConfigEnable(bool& value) { + std::string str = GetConfigStr(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_ENABLE, CONFIG_GPU_RESOURCE_ENABLE_DEFAULT); + Status s = CheckGpuResourceConfigEnable(str); if (!s.ok()) { return s; } @@ -1006,13 +1005,13 @@ Config::GetGpuResourceConfigEnableGpu(bool& value) { Status Config::GetGpuResourceConfigCacheCapacity(int64_t& value) { - bool enable_gpu = false; - Status s = GetGpuResourceConfigEnableGpu(enable_gpu); + bool gpu_resource_enable = false; + Status s = GetGpuResourceConfigEnable(gpu_resource_enable); if (!s.ok()) { return s; } - if (!enable_gpu) { - std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable_gpu is set to false."; + if (!gpu_resource_enable) { + std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable is set to false."; return Status(SERVER_UNSUPPORTED_ERROR, msg); } std::string str = GetConfigStr(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_CACHE_CAPACITY, @@ -1021,19 +1020,19 @@ Config::GetGpuResourceConfigCacheCapacity(int64_t& value) { if (!s.ok()) { return s; } - value = std::stoi(str); + value = std::stoll(str); return Status::OK(); } Status Config::GetGpuResourceConfigCacheThreshold(float& value) { - bool enable_gpu = false; - Status s = GetGpuResourceConfigEnableGpu(enable_gpu); + bool gpu_resource_enable = false; + Status s = GetGpuResourceConfigEnable(gpu_resource_enable); if (!s.ok()) { return s; } - if (!enable_gpu) { - std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable_gpu is set to false."; + if (!gpu_resource_enable) { + std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable is set to false."; return Status(SERVER_UNSUPPORTED_ERROR, msg); } std::string str = GetConfigStr(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_CACHE_THRESHOLD, @@ -1047,14 +1046,14 @@ Config::GetGpuResourceConfigCacheThreshold(float& value) { } Status -Config::GetGpuResourceConfigSearchResources(std::vector& value) { - bool enable_gpu = false; - Status s = GetGpuResourceConfigEnableGpu(enable_gpu); +Config::GetGpuResourceConfigSearchResources(std::vector& value) { + bool gpu_resource_enable = false; + Status s = GetGpuResourceConfigEnable(gpu_resource_enable); if (!s.ok()) { return s; } - if (!enable_gpu) { - std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable_gpu is set to false."; + if (!gpu_resource_enable) { + std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable is set to false."; return Status(SERVER_UNSUPPORTED_ERROR, msg); } std::string str = GetConfigSequenceStr(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_SEARCH_RESOURCES, @@ -1066,20 +1065,20 @@ Config::GetGpuResourceConfigSearchResources(std::vector& value) { return s; } for (std::string& res : res_vec) { - value.push_back(std::stoi(res.substr(3))); + value.push_back(std::stoll(res.substr(3))); } return Status::OK(); } Status -Config::GetGpuResourceConfigBuildIndexResources(std::vector& value) { - bool enable_gpu = false; - Status s = GetGpuResourceConfigEnableGpu(enable_gpu); +Config::GetGpuResourceConfigBuildIndexResources(std::vector& value) { + bool gpu_resource_enable = false; + Status s = GetGpuResourceConfigEnable(gpu_resource_enable); if (!s.ok()) { return s; } - if (!enable_gpu) { - std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable_gpu is set to false."; + if (!gpu_resource_enable) { + std::string msg = "GPU not supported. Possible reason: gpu_resource_config.enable is set to false."; return Status(SERVER_UNSUPPORTED_ERROR, msg); } std::string str = @@ -1092,7 +1091,7 @@ Config::GetGpuResourceConfigBuildIndexResources(std::vector& value) { return s; } for (std::string& res : res_vec) { - value.push_back(std::stoi(res.substr(3))); + value.push_back(std::stoll(res.substr(3))); } return Status::OK(); } @@ -1295,12 +1294,12 @@ Config::SetEngineConfigGpuSearchThreshold(const std::string& value) { /* gpu resource config */ Status -Config::SetGpuResourceConfigEnableGpu(const std::string& value) { - Status s = CheckGpuResourceConfigEnableGpu(value); +Config::SetGpuResourceConfigEnable(const std::string& value) { + Status s = CheckGpuResourceConfigEnable(value); if (!s.ok()) { return s; } - SetConfigValueInMem(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_ENABLE_GPU, value); + SetConfigValueInMem(CONFIG_GPU_RESOURCE, CONFIG_GPU_RESOURCE_ENABLE, value); return Status::OK(); } diff --git a/core/src/server/Config.h b/core/src/server/Config.h index 0442ae0626..0907080a6f 100644 --- a/core/src/server/Config.h +++ b/core/src/server/Config.h @@ -85,11 +85,11 @@ static const char* CONFIG_ENGINE_GPU_SEARCH_THRESHOLD_DEFAULT = "1000"; /* gpu resource config */ static const char* CONFIG_GPU_RESOURCE = "gpu_resource_config"; -static const char* CONFIG_GPU_RESOURCE_ENABLE_GPU = "enable_gpu"; +static const char* CONFIG_GPU_RESOURCE_ENABLE = "enable"; #ifdef MILVUS_GPU_VERSION -static const char* CONFIG_GPU_RESOURCE_ENABLE_GPU_DEFAULT = "true"; +static const char* CONFIG_GPU_RESOURCE_ENABLE_DEFAULT = "true"; #else -static const char* CONFIG_GPU_RESOURCE_ENABLE_GPU_DEFAULT = "false"; +static const char* CONFIG_GPU_RESOURCE_ENABLE_DEFAULT = "false"; #endif static const char* CONFIG_GPU_RESOURCE_CACHE_CAPACITY = "cache_capacity"; static const char* CONFIG_GPU_RESOURCE_CACHE_CAPACITY_DEFAULT = "4"; @@ -175,7 +175,7 @@ class Config { /* gpu resource config */ Status - CheckGpuResourceConfigEnableGpu(const std::string& value); + CheckGpuResourceConfigEnable(const std::string& value); Status CheckGpuResourceConfigCacheCapacity(const std::string& value); Status @@ -210,11 +210,11 @@ class Config { Status GetDBConfigBackendUrl(std::string& value); Status - GetDBConfigArchiveDiskThreshold(int32_t& value); + GetDBConfigArchiveDiskThreshold(int64_t& value); Status - GetDBConfigArchiveDaysThreshold(int32_t& value); + GetDBConfigArchiveDaysThreshold(int64_t& value); Status - GetDBConfigInsertBufferSize(int32_t& value); + GetDBConfigInsertBufferSize(int64_t& value); Status GetDBConfigPreloadTable(std::string& value); @@ -236,23 +236,23 @@ class Config { /* engine config */ Status - GetEngineConfigUseBlasThreshold(int32_t& value); + GetEngineConfigUseBlasThreshold(int64_t& value); Status - GetEngineConfigOmpThreadNum(int32_t& value); + GetEngineConfigOmpThreadNum(int64_t& value); Status - GetEngineConfigGpuSearchThreshold(int32_t& value); + GetEngineConfigGpuSearchThreshold(int64_t& value); /* gpu resource config */ Status - GetGpuResourceConfigEnableGpu(bool& value); + GetGpuResourceConfigEnable(bool& value); Status GetGpuResourceConfigCacheCapacity(int64_t& value); Status GetGpuResourceConfigCacheThreshold(float& value); Status - GetGpuResourceConfigSearchResources(std::vector& value); + GetGpuResourceConfigSearchResources(std::vector& value); Status - GetGpuResourceConfigBuildIndexResources(std::vector& value); + GetGpuResourceConfigBuildIndexResources(std::vector& value); public: /* server config */ @@ -305,7 +305,7 @@ class Config { /* gpu resource config */ Status - SetGpuResourceConfigEnableGpu(const std::string& value); + SetGpuResourceConfigEnable(const std::string& value); Status SetGpuResourceConfigCacheCapacity(const std::string& value); Status diff --git a/core/src/server/DBWrapper.cpp b/core/src/server/DBWrapper.cpp index 1ff914dc69..f7a43d671c 100644 --- a/core/src/server/DBWrapper.cpp +++ b/core/src/server/DBWrapper.cpp @@ -89,7 +89,7 @@ DBWrapper::StartService() { } // engine config - int32_t omp_thread; + int64_t omp_thread; s = config.GetEngineConfigOmpThreadNum(omp_thread); if (!s.ok()) { std::cerr << s.ToString() << std::endl; @@ -100,7 +100,7 @@ DBWrapper::StartService() { omp_set_num_threads(omp_thread); SERVER_LOG_DEBUG << "Specify openmp thread number: " << omp_thread; } else { - uint32_t sys_thread_cnt = 8; + int64_t sys_thread_cnt = 8; if (CommonUtil::GetSystemAvailableThreads(sys_thread_cnt)) { omp_thread = static_cast(ceil(sys_thread_cnt * 0.5)); omp_set_num_threads(omp_thread); @@ -108,7 +108,7 @@ DBWrapper::StartService() { } // init faiss global variable - int32_t use_blas_threshold; + int64_t use_blas_threshold; s = config.GetEngineConfigUseBlasThreshold(use_blas_threshold); if (!s.ok()) { std::cerr << s.ToString() << std::endl; @@ -119,7 +119,7 @@ DBWrapper::StartService() { // set archive config engine::ArchiveConf::CriteriaT criterial; - int32_t disk, days; + int64_t disk, days; s = config.GetDBConfigArchiveDiskThreshold(disk); if (!s.ok()) { std::cerr << s.ToString() << std::endl; diff --git a/core/src/server/grpc_impl/request/CreatePartitionRequest.cpp b/core/src/server/grpc_impl/request/CreatePartitionRequest.cpp index bc2e7dc105..3bd4a86ef6 100644 --- a/core/src/server/grpc_impl/request/CreatePartitionRequest.cpp +++ b/core/src/server/grpc_impl/request/CreatePartitionRequest.cpp @@ -51,7 +51,7 @@ CreatePartitionRequest::OnExecute() { return status; } - status = ValidationUtil::ValidateTableName(partition_param_->partition_name()); + status = ValidationUtil::ValidatePartitionName(partition_param_->partition_name()); if (!status.ok()) { return status; } diff --git a/core/src/server/grpc_impl/request/DropPartitionRequest.cpp b/core/src/server/grpc_impl/request/DropPartitionRequest.cpp index 217ca63a2a..0e29b6abe8 100644 --- a/core/src/server/grpc_impl/request/DropPartitionRequest.cpp +++ b/core/src/server/grpc_impl/request/DropPartitionRequest.cpp @@ -22,6 +22,7 @@ #include "utils/ValidationUtil.h" #include +#include namespace milvus { namespace server { @@ -38,23 +39,40 @@ DropPartitionRequest::Create(const ::milvus::grpc::PartitionParam* partition_par Status DropPartitionRequest::OnExecute() { - if (!partition_param_->partition_name().empty()) { - auto status = ValidationUtil::ValidateTableName(partition_param_->partition_name()); - if (!status.ok()) { - return status; - } - return DBWrapper::DB()->DropPartition(partition_param_->partition_name()); - } else { - auto status = ValidationUtil::ValidateTableName(partition_param_->table_name()); + std::string table_name = partition_param_->table_name(); + std::string partition_name = partition_param_->partition_name(); + std::string partition_tag = partition_param_->tag(); + if (!partition_name.empty()) { + auto status = ValidationUtil::ValidateTableName(partition_name); if (!status.ok()) { return status; } - status = ValidationUtil::ValidatePartitionTags({partition_param_->tag()}); + // check partition existence + engine::meta::TableSchema table_info; + table_info.table_id_ = partition_name; + status = DBWrapper::DB()->DescribeTable(table_info); + if (!status.ok()) { + if (status.code() == DB_NOT_FOUND) { + return Status(SERVER_TABLE_NOT_EXIST, + "Table " + table_name + "'s partition " + partition_name + " not found"); + } else { + return status; + } + } + + return DBWrapper::DB()->DropPartition(partition_name); + } else { + auto status = ValidationUtil::ValidateTableName(table_name); if (!status.ok()) { return status; } - return DBWrapper::DB()->DropPartitionByTag(partition_param_->table_name(), partition_param_->tag()); + + status = ValidationUtil::ValidatePartitionTags({partition_tag}); + if (!status.ok()) { + return status; + } + return DBWrapper::DB()->DropPartitionByTag(table_name, partition_tag); } } diff --git a/core/src/utils/CommonUtil.cpp b/core/src/utils/CommonUtil.cpp index fbf3112aeb..26e43619fb 100644 --- a/core/src/utils/CommonUtil.cpp +++ b/core/src/utils/CommonUtil.cpp @@ -54,7 +54,7 @@ CommonUtil::GetSystemMemInfo(uint64_t& total_mem, uint64_t& free_mem) { } bool -CommonUtil::GetSystemAvailableThreads(uint32_t& thread_count) { +CommonUtil::GetSystemAvailableThreads(int64_t& thread_count) { // threadCnt = std::thread::hardware_concurrency(); thread_count = sysconf(_SC_NPROCESSORS_CONF); thread_count *= THREAD_MULTIPLY_CPU; diff --git a/core/src/utils/CommonUtil.h b/core/src/utils/CommonUtil.h index 939bdd6d31..121196986a 100644 --- a/core/src/utils/CommonUtil.h +++ b/core/src/utils/CommonUtil.h @@ -30,7 +30,7 @@ class CommonUtil { static bool GetSystemMemInfo(uint64_t& total_mem, uint64_t& free_mem); static bool - GetSystemAvailableThreads(uint32_t& thread_count); + GetSystemAvailableThreads(int64_t& thread_count); static bool IsFileExist(const std::string& path); diff --git a/core/src/utils/ValidationUtil.cpp b/core/src/utils/ValidationUtil.cpp index 080de77e17..12b2372fc5 100644 --- a/core/src/utils/ValidationUtil.cpp +++ b/core/src/utils/ValidationUtil.cpp @@ -18,6 +18,7 @@ #include "utils/ValidationUtil.h" #include "Log.h" #include "db/engine/ExecutionEngine.h" +#include "utils/StringHelpFunctions.h" #include #ifdef MILVUS_GPU_VERSION @@ -168,11 +169,26 @@ ValidationUtil::ValidateSearchNprobe(int64_t nprobe, const engine::meta::TableSc return Status::OK(); } +Status +ValidationUtil::ValidatePartitionName(const std::string& partition_name) { + if (partition_name.empty()) { + std::string msg = "Partition name should not be empty."; + SERVER_LOG_ERROR << msg; + return Status(SERVER_INVALID_TABLE_NAME, msg); + } + + return ValidateTableName(partition_name); +} + Status ValidationUtil::ValidatePartitionTags(const std::vector& partition_tags) { - for (auto& tag : partition_tags) { - if (tag.empty()) { - std::string msg = "Invalid partition tag: " + tag + ". " + "Partition tag should not be empty."; + for (const std::string& tag : partition_tags) { + // trim side-blank of tag, only compare valid characters + // for example: " ab cd " is treated as "ab cd" + std::string valid_tag = tag; + StringHelpFunctions::TrimStringBlank(valid_tag); + if (valid_tag.empty()) { + std::string msg = "Invalid partition tag: " + valid_tag + ". " + "Partition tag should not be empty."; SERVER_LOG_ERROR << msg; return Status(SERVER_INVALID_NPROBE, msg); } diff --git a/core/src/utils/ValidationUtil.h b/core/src/utils/ValidationUtil.h index 201ccef3bd..ab32c35c40 100644 --- a/core/src/utils/ValidationUtil.h +++ b/core/src/utils/ValidationUtil.h @@ -55,6 +55,9 @@ class ValidationUtil { static Status ValidateSearchNprobe(int64_t nprobe, const engine::meta::TableSchema& table_schema); + static Status + ValidatePartitionName(const std::string& partition_name); + static Status ValidatePartitionTags(const std::vector& partition_tags); diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/wrapper/ConfAdapter.cpp index 0ab38d3394..aa4b3c12b8 100644 --- a/core/src/wrapper/ConfAdapter.cpp +++ b/core/src/wrapper/ConfAdapter.cpp @@ -201,5 +201,35 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) return conf; } +knowhere::Config +SPTAGKDTConfAdapter::Match(const TempMetaConf& metaconf) { + auto conf = std::make_shared(); + conf->d = metaconf.dim; + conf->metric_type = metaconf.metric_type; + return conf; +} + +knowhere::Config +SPTAGKDTConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) { + auto conf = std::make_shared(); + conf->k = metaconf.k; + return conf; +} + +knowhere::Config +SPTAGBKTConfAdapter::Match(const TempMetaConf& metaconf) { + auto conf = std::make_shared(); + conf->d = metaconf.dim; + conf->metric_type = metaconf.metric_type; + return conf; +} + +knowhere::Config +SPTAGBKTConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) { + auto conf = std::make_shared(); + conf->k = metaconf.k; + return conf; +} + } // namespace engine } // namespace milvus diff --git a/core/src/wrapper/ConfAdapter.h b/core/src/wrapper/ConfAdapter.h index 85637a4969..46fc27eb3b 100644 --- a/core/src/wrapper/ConfAdapter.h +++ b/core/src/wrapper/ConfAdapter.h @@ -94,5 +94,23 @@ class NSGConfAdapter : public IVFConfAdapter { MatchSearch(const TempMetaConf& metaconf, const IndexType& type) final; }; +class SPTAGKDTConfAdapter : public ConfAdapter { + public: + knowhere::Config + Match(const TempMetaConf& metaconf) override; + + knowhere::Config + MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override; +}; + +class SPTAGBKTConfAdapter : public ConfAdapter { + public: + knowhere::Config + Match(const TempMetaConf& metaconf) override; + + knowhere::Config + MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override; +}; + } // namespace engine } // namespace milvus diff --git a/core/src/wrapper/ConfAdapterMgr.cpp b/core/src/wrapper/ConfAdapterMgr.cpp index d0eba04529..cf58c0110f 100644 --- a/core/src/wrapper/ConfAdapterMgr.cpp +++ b/core/src/wrapper/ConfAdapterMgr.cpp @@ -56,6 +56,9 @@ AdapterMgr::RegisterAdapter() { REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix); REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix); + + REGISTER_CONF_ADAPTER(SPTAGKDTConfAdapter, IndexType::SPTAG_KDT_RNT_CPU, sptag_kdt); + REGISTER_CONF_ADAPTER(SPTAGBKTConfAdapter, IndexType::SPTAG_BKT_RNT_CPU, sptag_bkt); } } // namespace engine diff --git a/core/src/wrapper/KnowhereResource.cpp b/core/src/wrapper/KnowhereResource.cpp index bf6675cafa..5a2296b16e 100644 --- a/core/src/wrapper/KnowhereResource.cpp +++ b/core/src/wrapper/KnowhereResource.cpp @@ -48,7 +48,7 @@ KnowhereResource::Initialize() { // get build index gpu resource server::Config& config = server::Config::GetInstance(); - std::vector build_index_gpus; + std::vector build_index_gpus; s = config.GetGpuResourceConfigBuildIndexResources(build_index_gpus); if (!s.ok()) return s; @@ -58,7 +58,7 @@ KnowhereResource::Initialize() { } // get search gpu resource - std::vector search_gpus; + std::vector search_gpus; s = config.GetGpuResourceConfigSearchResources(search_gpus); if (!s.ok()) return s; diff --git a/core/src/wrapper/VecImpl.cpp b/core/src/wrapper/VecImpl.cpp index e7967cbf59..74e9e94a2f 100644 --- a/core/src/wrapper/VecImpl.cpp +++ b/core/src/wrapper/VecImpl.cpp @@ -21,17 +21,19 @@ #include "knowhere/index/vector_index/IndexIDMAP.h" #include "utils/Log.h" #include "wrapper/WrapperException.h" +#include "wrapper/gpu/GPUVecImpl.h" #ifdef MILVUS_GPU_VERSION -#include -#include +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" #endif /* * no parameter check in this layer. - * only responible for index combination + * only responsible for index combination */ namespace milvus { diff --git a/core/src/wrapper/VecIndex.cpp b/core/src/wrapper/VecIndex.cpp index 81a13c60f9..75f75fb983 100644 --- a/core/src/wrapper/VecIndex.cpp +++ b/core/src/wrapper/VecIndex.cpp @@ -22,8 +22,8 @@ #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" -#include "knowhere/index/vector_index/IndexKDT.h" #include "knowhere/index/vector_index/IndexNSG.h" +#include "knowhere/index/vector_index/IndexSPTAG.h" #include "utils/Log.h" #ifdef MILVUS_GPU_VERSION @@ -128,7 +128,11 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) { break; } case IndexType::SPTAG_KDT_RNT_CPU: { - index = std::make_shared(); + index = std::make_shared("KDT"); + break; + } + case IndexType::SPTAG_BKT_RNT_CPU: { + index = std::make_shared("BKT"); break; } case IndexType::FAISS_IVFSQ8_CPU: { diff --git a/core/src/wrapper/VecIndex.h b/core/src/wrapper/VecIndex.h index efe01a25d7..e69655b087 100644 --- a/core/src/wrapper/VecIndex.h +++ b/core/src/wrapper/VecIndex.h @@ -49,6 +49,7 @@ enum class IndexType { FAISS_IVFSQ8_HYBRID, // only support build on gpu. NSG_MIX, FAISS_IVFPQ_MIX, + SPTAG_BKT_RNT_CPU, }; class VecIndex; @@ -139,6 +140,9 @@ write_index(VecIndexPtr index, const std::string& location); extern VecIndexPtr read_index(const std::string& location); +VecIndexPtr +read_index(const std::string& location, knowhere::BinarySet& index_binary); + extern VecIndexPtr GetVecIndexFactory(const IndexType& type, const Config& cfg = Config()); diff --git a/core/ubuntu_build_deps.sh b/core/ubuntu_build_deps.sh index 911046aa1f..6f3fabc084 100755 --- a/core/ubuntu_build_deps.sh +++ b/core/ubuntu_build_deps.sh @@ -9,4 +9,6 @@ sudo apt-get -y update && sudo apt-get -y install intel-mkl-gnu-2019.5-281 intel sudo apt-get install -y gfortran libmysqlclient-dev mysql-client libcurl4-openssl-dev libboost-system-dev \ libboost-filesystem-dev libboost-serialization-dev libboost-regex-dev -sudo ln -s /usr/lib/x86_64-linux-gnu/libmysqlclient.so /usr/lib/x86_64-linux-gnu/libmysqlclient_r.so +if [ ! -f "/usr/lib/x86_64-linux-gnu/libmysqlclient_r.so" ]; then + sudo ln -s /usr/lib/x86_64-linux-gnu/libmysqlclient.so /usr/lib/x86_64-linux-gnu/libmysqlclient_r.so +fi diff --git a/core/unittest/db/test_db.cpp b/core/unittest/db/test_db.cpp index 9e01b5ea1a..217fbe429e 100644 --- a/core/unittest/db/test_db.cpp +++ b/core/unittest/db/test_db.cpp @@ -467,6 +467,13 @@ TEST_F(DBTest, PARTITION_TEST) { stat = db_->CreatePartition(table_name, partition_name, partition_tag); ASSERT_TRUE(stat.ok()); + // not allow nested partition + stat = db_->CreatePartition(partition_name, "dumy", "dummy"); + ASSERT_FALSE(stat.ok()); + + // not allow duplicated partition + stat = db_->CreatePartition(table_name, partition_name, partition_tag); + ASSERT_FALSE(stat.ok()); std::vector xb; BuildVectors(INSERT_BATCH, xb); diff --git a/core/unittest/db/test_db_mysql.cpp b/core/unittest/db/test_db_mysql.cpp index f828431838..93eea2fd27 100644 --- a/core/unittest/db/test_db_mysql.cpp +++ b/core/unittest/db/test_db_mysql.cpp @@ -297,6 +297,14 @@ TEST_F(MySqlDBTest, PARTITION_TEST) { stat = db_->CreatePartition(table_name, partition_name, partition_tag); ASSERT_TRUE(stat.ok()); + // not allow nested partition + stat = db_->CreatePartition(partition_name, "dumy", "dummy"); + ASSERT_FALSE(stat.ok()); + + // not allow duplicated partition + stat = db_->CreatePartition(table_name, partition_name, partition_tag); + ASSERT_FALSE(stat.ok()); + std::vector xb; BuildVectors(INSERT_BATCH, xb); diff --git a/core/unittest/server/test_config.cpp b/core/unittest/server/test_config.cpp index 6ba7f19160..664a08d631 100644 --- a/core/unittest/server/test_config.cpp +++ b/core/unittest/server/test_config.cpp @@ -104,7 +104,6 @@ TEST_F(ConfigTest, SERVER_CONFIG_VALID_TEST) { milvus::server::Config& config = milvus::server::Config::GetInstance(); milvus::Status s; std::string str_val; - int32_t int32_val; int64_t int64_val; float float_val; bool bool_val; @@ -160,26 +159,26 @@ TEST_F(ConfigTest, SERVER_CONFIG_VALID_TEST) { ASSERT_TRUE(s.ok()); ASSERT_TRUE(str_val == db_backend_url); - int32_t db_archive_disk_threshold = 100; + int64_t db_archive_disk_threshold = 100; s = config.SetDBConfigArchiveDiskThreshold(std::to_string(db_archive_disk_threshold)); ASSERT_TRUE(s.ok()); - s = config.GetDBConfigArchiveDiskThreshold(int32_val); + s = config.GetDBConfigArchiveDiskThreshold(int64_val); ASSERT_TRUE(s.ok()); - ASSERT_TRUE(int32_val == db_archive_disk_threshold); + ASSERT_TRUE(int64_val == db_archive_disk_threshold); - int32_t db_archive_days_threshold = 365; + int64_t db_archive_days_threshold = 365; s = config.SetDBConfigArchiveDaysThreshold(std::to_string(db_archive_days_threshold)); ASSERT_TRUE(s.ok()); - s = config.GetDBConfigArchiveDaysThreshold(int32_val); + s = config.GetDBConfigArchiveDaysThreshold(int64_val); ASSERT_TRUE(s.ok()); - ASSERT_TRUE(int32_val == db_archive_days_threshold); + ASSERT_TRUE(int64_val == db_archive_days_threshold); - int32_t db_insert_buffer_size = 2; + int64_t db_insert_buffer_size = 2; s = config.SetDBConfigInsertBufferSize(std::to_string(db_insert_buffer_size)); ASSERT_TRUE(s.ok()); - s = config.GetDBConfigInsertBufferSize(int32_val); + s = config.GetDBConfigInsertBufferSize(int64_val); ASSERT_TRUE(s.ok()); - ASSERT_TRUE(int32_val == db_insert_buffer_size); + ASSERT_TRUE(int64_val == db_insert_buffer_size); /* metric config */ bool metric_enable_monitor = false; @@ -223,32 +222,32 @@ TEST_F(ConfigTest, SERVER_CONFIG_VALID_TEST) { ASSERT_TRUE(bool_val == cache_insert_data); /* engine config */ - int32_t engine_use_blas_threshold = 50; + int64_t engine_use_blas_threshold = 50; s = config.SetEngineConfigUseBlasThreshold(std::to_string(engine_use_blas_threshold)); ASSERT_TRUE(s.ok()); - s = config.GetEngineConfigUseBlasThreshold(int32_val); + s = config.GetEngineConfigUseBlasThreshold(int64_val); ASSERT_TRUE(s.ok()); - ASSERT_TRUE(int32_val == engine_use_blas_threshold); + ASSERT_TRUE(int64_val == engine_use_blas_threshold); - int32_t engine_omp_thread_num = 8; + int64_t engine_omp_thread_num = 8; s = config.SetEngineConfigOmpThreadNum(std::to_string(engine_omp_thread_num)); ASSERT_TRUE(s.ok()); - s = config.GetEngineConfigOmpThreadNum(int32_val); + s = config.GetEngineConfigOmpThreadNum(int64_val); ASSERT_TRUE(s.ok()); - ASSERT_TRUE(int32_val == engine_omp_thread_num); + ASSERT_TRUE(int64_val == engine_omp_thread_num); - int32_t engine_gpu_search_threshold = 800; + int64_t engine_gpu_search_threshold = 800; s = config.SetEngineConfigGpuSearchThreshold(std::to_string(engine_gpu_search_threshold)); ASSERT_TRUE(s.ok()); - s = config.GetEngineConfigGpuSearchThreshold(int32_val); + s = config.GetEngineConfigGpuSearchThreshold(int64_val); ASSERT_TRUE(s.ok()); - ASSERT_TRUE(int32_val == engine_gpu_search_threshold); + ASSERT_TRUE(int64_val == engine_gpu_search_threshold); /* gpu resource config */ bool resource_enable_gpu = true; - s = config.SetGpuResourceConfigEnableGpu(std::to_string(resource_enable_gpu)); + s = config.SetGpuResourceConfigEnable(std::to_string(resource_enable_gpu)); ASSERT_TRUE(s.ok()); - s = config.GetGpuResourceConfigEnableGpu(bool_val); + s = config.GetGpuResourceConfigEnable(bool_val); ASSERT_TRUE(s.ok()); ASSERT_TRUE(bool_val == resource_enable_gpu); @@ -267,7 +266,7 @@ TEST_F(ConfigTest, SERVER_CONFIG_VALID_TEST) { ASSERT_TRUE(float_val == gpu_cache_threshold); std::vector search_resources = {"gpu0"}; - std::vector search_res_vec; + std::vector search_res_vec; std::string search_res_str; milvus::server::StringHelpFunctions::MergeStringWithDelimeter( search_resources, milvus::server::CONFIG_GPU_RESOURCE_DELIMITER, search_res_str); @@ -276,11 +275,11 @@ TEST_F(ConfigTest, SERVER_CONFIG_VALID_TEST) { s = config.GetGpuResourceConfigSearchResources(search_res_vec); ASSERT_TRUE(s.ok()); for (size_t i = 0; i < search_resources.size(); i++) { - ASSERT_TRUE(std::stoi(search_resources[i].substr(3)) == search_res_vec[i]); + ASSERT_TRUE(std::stoll(search_resources[i].substr(3)) == search_res_vec[i]); } std::vector build_index_resources = {"gpu0"}; - std::vector build_index_res_vec; + std::vector build_index_res_vec; std::string build_index_res_str; milvus::server::StringHelpFunctions::MergeStringWithDelimeter( build_index_resources, milvus::server::CONFIG_GPU_RESOURCE_DELIMITER, build_index_res_str); @@ -289,7 +288,7 @@ TEST_F(ConfigTest, SERVER_CONFIG_VALID_TEST) { s = config.GetGpuResourceConfigBuildIndexResources(build_index_res_vec); ASSERT_TRUE(s.ok()); for (size_t i = 0; i < build_index_resources.size(); i++) { - ASSERT_TRUE(std::stoi(build_index_resources[i].substr(3)) == build_index_res_vec[i]); + ASSERT_TRUE(std::stoll(build_index_resources[i].substr(3)) == build_index_res_vec[i]); } #endif } @@ -394,7 +393,7 @@ TEST_F(ConfigTest, SERVER_CONFIG_INVALID_TEST) { ASSERT_FALSE(s.ok()); /* gpu resource config */ - s = config.SetGpuResourceConfigEnableGpu("ok"); + s = config.SetGpuResourceConfigEnable("ok"); ASSERT_FALSE(s.ok()); #ifdef MILVUS_GPU_VERSION diff --git a/core/unittest/server/test_rpc.cpp b/core/unittest/server/test_rpc.cpp index 34d4b0ad3a..5753c68422 100644 --- a/core/unittest/server/test_rpc.cpp +++ b/core/unittest/server/test_rpc.cpp @@ -409,7 +409,7 @@ TEST_F(RpcHandlerTest, PARTITION_TEST) { partition_parm.set_partition_name(partition_name); handler->DropPartition(&context, &partition_parm, &response); - ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code()); + ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code()); } TEST_F(RpcHandlerTest, CMD_TEST) { diff --git a/core/unittest/server/test_util.cpp b/core/unittest/server/test_util.cpp index 68400f2454..e5884cac65 100644 --- a/core/unittest/server/test_util.cpp +++ b/core/unittest/server/test_util.cpp @@ -60,7 +60,7 @@ TEST(UtilTest, COMMON_TEST) { ASSERT_GT(total_mem, 0); ASSERT_GT(free_mem, 0); - uint32_t thread_cnt = 0; + int64_t thread_cnt = 0; milvus::server::CommonUtil::GetSystemAvailableThreads(thread_cnt); ASSERT_GT(thread_cnt, 0); diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp index a07fafc7b8..4019c0f63c 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/unittest/wrapper/test_wrapper.cpp @@ -16,28 +16,29 @@ // under the License. #include "easyloggingpp/easylogging++.h" -#include "wrapper/VecIndex.h" #ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif #include "knowhere/index/vector_index/helpers/IndexParameter.h" +#include "wrapper/VecIndex.h" #include "wrapper/utils.h" #include INITIALIZE_EASYLOGGINGPP +using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; -using ::testing::Combine; class KnowhereWrapperTest - : public DataGenBase, - public TestWithParam<::std::tuple> { + : public DataGenBase, + public TestWithParam<::std::tuple> { protected: - void SetUp() override { + void + SetUp() override { #ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); #endif @@ -58,7 +59,8 @@ class KnowhereWrapperTest searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf); } - void TearDown() override { + void + TearDown() override { #ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif @@ -71,22 +73,20 @@ class KnowhereWrapperTest knowhere::Config searchconf; }; -INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, - Values( - //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] - +INSTANTIATE_TEST_CASE_P( + WrapperParam, KnowhereWrapperTest, + Values( + //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] #ifdef MILVUS_GPU_VERSION std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 1000, 10, 10), - // std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, - // 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFPQ_MIX, "Default", 64, 1000, 10, 10), - -// std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10), + // std::make_tuple(milvus::engine::IndexType::NSG_MIX, "Default", 128, 250000, 10, 10), #endif - // std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10), + // std::make_tuple(milvus::engine::IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 100, 10, 10), + // std::make_tuple(milvus::engine::IndexType::SPTAG_BKT_RNT_CPU, "Default", 128, 100, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 1000, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 1000, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10))); diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000000..f5f964de0d --- /dev/null +++ b/docs/README.md @@ -0,0 +1,23 @@ +# Docs + +This repository contains test reports on the search performance of different index types on standalone Milvus. + +The tests are run on [SIFT1B dataset](http://corpus-texmex.irisa.fr/), and provide results on the following measures: + +- Query Elapsed Time: Time cost (in seconds) to run a query. + +- Recall: The fraction of the total amount of relevant instances that were actually retrieved. + +Test variables are `nq` and `topk`. + +## Test reports + +The following is a list of existing test reports: + +- [IVF_SQ8](test_report/milvus_ivfsq8_test_report_detailed_version.md) +- [IVF_SQ8H](test_report/milvus_ivfsq8h_test_report_detailed_version.md) + +To read the CN version of these reports: + +- [IVF_SQ8_cn](test_report/milvus_ivfsq8_test_report_detailed_version_cn.md) +- [IVF_SQ8H_cn](test_report/milvus_ivfsq8h_test_report_detailed_version_cn.md) diff --git a/docs/test_report/milvus_ivfsq8h_test_report_detailed_version.md b/docs/test_report/milvus_ivfsq8h_test_report_detailed_version.md new file mode 100755 index 0000000000..042d059684 --- /dev/null +++ b/docs/test_report/milvus_ivfsq8h_test_report_detailed_version.md @@ -0,0 +1,179 @@ +# milvus_ivfsq8h_test_report_detailed_version + +## Summary + +This document contains the test reports of IVF_SQ8H index on Milvus single server. + + + +## Test objectives + +The time cost and recall when searching with different parameters. + + + +## Test method + +### Hardware/Software requirements + +Operating System: CentOS Linux release 7.6.1810 (Core) + +CPU: Intel(R) Xeon(R) CPU E5-2678 v3 @ 2.50GHz + +GPU0: GeForce GTX 1080 + +GPU1: GeForce GTX 1080 + +Memory: 503GB + +Docker version: 18.09 + +NVIDIA Driver version: 430.34 + +Milvus version: 0.5.3 + +SDK interface: Python 3.6.8 + +pymilvus version: 0.2.5 + + + +### Data model + +The data used in the tests are: + +- Data source: sift1b +- Data type: hdf5 + +For details on this dataset, please check : http://corpus-texmex.irisa.fr/ . + + + +### Measures + +- Query Elapsed Time: Time cost (in seconds) to run a query. Variables that affect Query Elapsed Time: + + - nq (Number of queried vectors) + +> Note: In the query test of query elapsed time, we will test the following parameters with different values: + > + > nq - grouped by: [1, 5, 10, 50, 100, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800]. + > + +- Recall: The fraction of the total amount of relevant instances that were actually retrieved . Variables that affect Recall: + + - nq (Number of queried vectors) + - topk (Top k result of a query) + + > Note: In the query test of recall, we will test the following parameters with different values: + > + > nq - grouped by: [10, 50, 100, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800], + > + > topk - grouped by: [1, 10, 100] + + + +## Test reports + +### Test environment + +Data base: sift1b-1,000,000,000 vectors, 128-dimension + +Table Attributes + +- nlist: 16384 +- metric_type: L2 + +Query configuration + +- nprobe: 32 + +Milvus configuration + +- cpu_cache_capacity: 150 +- gpu_cache_capacity: 6 +- use_blas_threshold: 1100 +- gpu_search_threshold: 1200 +- search_resources: cpu, gpu0, gpu1 + +The definitions of Milvus configuration are on https://milvus.io/docs/en/reference/milvus_config/. + +Test method + +Test the query elapsed time and recall with several parameters, and once only change one parameter. + +- Whether to restart Milvus after each query: No + + + +### Performance test + +#### Data query + +**Test result** + +Query Elapsed Time + +topk = 100 + +| nq/topk | topk=100 | +| :-----: | :------: | +| nq=1 | 0.34 | +| nq=5 | 0.72 | +| nq=10 | 0.91 | +| nq=50 | 1.51 | +| nq=100 | 2.49 | +| nq=200 | 4.09 | +| nq=400 | 7.32 | +| nq=600 | 10.63 | +| nq=800 | 13.84 | +| nq=1000 | 16.83 | +| nq=1200 | 18.20 | +| nq=1400 | 20.1 | +| nq=1600 | 20.0 | +| nq=1800 | 19.86 | + +When nq is 1800, the query time cost of a 128-dimension vector is around 11ms. + + + +**Conclusion** + +When nq < 1200, the query elapsed time increases quickly with nq; when nq > 1200, the query elapsed time increases much slower. It is because gpu_search_threshold is set to 1200, when nq < 1200, CPU is chosen to do the query, otherwise GPU is chosen. Compared with CPU, GPU has much more cores and stronger computing capability. When nq is large, it can better reflect GPU's advantages on computing. + +The query elapsed time consists of two parts: (1) index CPU-to-GPU copy time; (2) nprobe buckets search time. When nq is larger enough, index CPU-to-GPU copy time can be amortized efficiently. So Milvus performs well through setting suitable gpu_search_threshold. + + + +### Recall test + +**Test result** + +topk = 1 : recall - recall@1 + +topk = 10 : recall - recall@10 + +topk = 100 : recall - recall@100 + +We use the ground_truth in sift1b dataset to calculate the recall of query results. + +| nq/topk | topk=1 | topk=10 | topk=100 | +| :-----: | :----: | :-----: | :------: | +| nq=10 | 0.900 | 0.910 | 0.939 | +| nq=50 | 0.980 | 0.950 | 0.941 | +| nq=100 | 0.970 | 0.937 | 0.931 | +| nq=200 | 0.955 | 0.941 | 0.929 | +| nq=400 | 0.958 | 0.944 | 0.932 | +| nq=600 | 0.952 | 0.946 | 0.934 | +| nq=800 | 0.941 | 0.943 | 0.930 | +| nq=1000 | 0.938 | 0.942 | 0.930 | +| nq=1200 | 0.937 | 0.943 | 0.931 | +| nq=1400 | 0.939 | 0.945 | 0.931 | +| nq=1600 | 0.936 | 0.945 | 0.931 | +| nq=1800 | 0.937 | 0.946 | 0.932 | + + + +**Conclusion** + +As nq increases, the recall gradually stabilizes to over 93%. The usage of CPU or GPU and different topk are not related to recall. \ No newline at end of file diff --git a/docs/test_report/milvus_ivfsq8h_test_report_detailed_version_cn.md b/docs/test_report/milvus_ivfsq8h_test_report_detailed_version_cn.md new file mode 100644 index 0000000000..b50d00f9bd --- /dev/null +++ b/docs/test_report/milvus_ivfsq8h_test_report_detailed_version_cn.md @@ -0,0 +1,180 @@ +# milvus_ivfsq8h_test_report_detailed_version_cn + +## 概述 + +本文描述了ivfsq8h索引在milvus单机部署方式下的测试结果。 + + + +## 测试目标 + +参数不同情况下的查询时间和召回率。 + + + +## 测试方法 + +### 软硬件环境 + +操作系统：CentOS Linux release 7.6.1810 (Core) + +CPU：Intel(R) Xeon(R) CPU E5-2678 v3 @ 2.50GHz + +GPU0：GeForce GTX 1080 + +GPU1：GeForce GTX 1080 + +内存：503GB + +Docker版本：18.09 + +NVIDIA Driver版本：430.34 + +Milvus版本：0.5.3 + +SDK接口：Python 3.6.8 + +pymilvus版本：0.2.5 + + + +### 数据模型 + +本测试中用到的主要数据: + +- 数据来源：sift1b +- 数据类型：hdf5 + +关于该数据集的详细信息请参考：http://corpus-texmex.irisa.fr/ 。 + + + +### 测试指标 + +- Query Elapsed Time: 数据库查询所有向量的时间（以秒计）。影响Query Elapsed Time的变量： + + - nq (被查询向量的数量) + +> 备注：在向量查询测试中，我们会测试下面参数不同的取值来观察结果： + > + > 被查询向量的数量nq将按照 [1, 5, 10, 50, 100, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800]的数量分组。 + > + +- Recall: 实际返回的正确结果占总数之比。影响Recall的变量： + + - nq (被查询向量的数量) + - topk (单条查询中最相似的K个结果) + + > 备注：在向量准确性测试中，我们会测试下面参数不同的取值来观察结果： + > + > 被查询向量的数量nq将按照 [10, 50, 100, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800]的数量分组， + > + > 单条查询中最相似的K个结果topk将按照[1, 10, 100]的数量分组。 + + + +## 测试报告 + +### 测试环境 + +数据集: sift1b-1,000,000,000向量, 128维 + +表格属性： + +- nlist: 16384 +- metric_type: L2 + +查询设置： + +- nprobe: 32 + +Milvus设置： + +- cpu_cache_capacity: 150 +- gpu_cache_capacity: 6 +- use_blas_threshold: 1100 +- gpu_search_threshold: 1200 +- search_resources: cpu, gpu0, gpu1 + +Milvus设置的详细定义可以参考 https://milvus.io/docs/en/reference/milvus_config/ 。 + +测试方法 + +通过一次仅改变一个参数的值，测试查询向量时间和召回率。 + +- 查询后是否重启Milvus：否 + + + +### 性能测试 + +#### 数据查询 + +测试结果 + +Query Elapsed Time + +topk = 100 + +| nq/topk | topk=100 | +| :-----: | :------: | +| nq=1 | 0.34 | +| nq=5 | 0.72 | +| nq=10 | 0.91 | +| nq=50 | 1.51 | +| nq=100 | 2.49 | +| nq=200 | 4.09 | +| nq=400 | 7.32 | +| nq=600 | 10.63 | +| nq=800 | 13.84 | +| nq=1000 | 16.83 | +| nq=1200 | 18.20 | +| nq=1400 | 20.1 | +| nq=1600 | 20.0 | +| nq=1800 | 19.86 | + +当nq为1800时，查询一条128维向量需要耗时约11毫秒。 + + + +**总结** + +当nq小于1200时，查询耗时随nq的增长快速增大；当nq大于1200时，查询耗时的增大则缓慢许多。这是因为gpu_search_threshold这一参数的值被设为1200，当nq<1200时，选择CPU进行操作，否则选择GPU进行操作。与CPU。 + +在GPU模式下的查询耗时由两部分组成：（1）索引从CPU到GPU的拷贝时间；（2）所有分桶的查询时间。当nq小于500时，索引从CPU到GPU 的拷贝时间无法被有效均摊，此时CPU模式时一个更优的选择；当nq大于500时，选择GPU模式更合理。和CPU相比，GPU具有更多的核数和更强的算力。当nq较大时，GPU在计算上的优势能被更好地被体现。 + + + +### 召回率测试 + +**测试结果** + +topk = 1 : recall - recall@1 + +topk = 10 : recall - recall@10 + +topk = 100 : recall - recall@100 + +我们利用sift1b数据集中的ground_truth来计算查询结果的召回率。 + +| nq/topk | topk=1 | topk=10 | topk=100 | +| :-----: | :----: | :-----: | :------: | +| nq=10 | 0.900 | 0.910 | 0.939 | +| nq=50 | 0.980 | 0.950 | 0.941 | +| nq=100 | 0.970 | 0.937 | 0.931 | +| nq=200 | 0.955 | 0.941 | 0.929 | +| nq=400 | 0.958 | 0.944 | 0.932 | +| nq=600 | 0.952 | 0.946 | 0.934 | +| nq=800 | 0.941 | 0.943 | 0.930 | +| nq=1000 | 0.938 | 0.942 | 0.930 | +| nq=1200 | 0.937 | 0.943 | 0.931 | +| nq=1400 | 0.939 | 0.945 | 0.931 | +| nq=1600 | 0.936 | 0.945 | 0.931 | +| nq=1800 | 0.937 | 0.946 | 0.932 | + + + +**总结** + +随着nq的增大，召回率逐渐稳定至93%以上。CPU/GPU的使用以及topk的值与召回率的大小无关。 + diff --git a/shards/mishards/connections.py b/shards/mishards/connections.py index 618690a099..50e214ec9a 100644 --- a/shards/mishards/connections.py +++ b/shards/mishards/connections.py @@ -2,6 +2,7 @@ import logging import threading from functools import wraps from milvus import Milvus +from milvus.client.hooks import BaseaSearchHook from mishards import (settings, exceptions) from utils import singleton @@ -9,6 +10,12 @@ from utils import singleton logger = logging.getLogger(__name__) +class Searchook(BaseaSearchHook): + + def on_response(self, *args, **kwargs): + return True + + class Connection: def __init__(self, name, uri, max_retry=1, error_handlers=None, **kwargs): self.name = name @@ -18,6 +25,9 @@ class Connection: self.conn = Milvus() self.error_handlers = [] if not error_handlers else error_handlers self.on_retry_func = kwargs.get('on_retry_func', None) + + # define search hook + self.conn._set_hook(search_in_file=Searchook()) # self._connect() def __str__(self): diff --git a/shards/mishards/service_handler.py b/shards/mishards/service_handler.py index 2f19152ae6..fc0ee0fa2b 100644 --- a/shards/mishards/service_handler.py +++ b/shards/mishards/service_handler.py @@ -29,39 +29,71 @@ class ServiceHandler(milvus_pb2_grpc.MilvusServiceServicer): self.router = router self.max_workers = max_workers + def _reduce(self, source_ids, ids, source_diss, diss, k, reverse): + if source_diss[k - 1] <= diss[0]: + return source_ids, source_diss + if diss[k - 1] <= source_diss[0]: + return ids, diss + + source_diss.extend(diss) + diss_t = enumerate(source_diss) + diss_m_rst = sorted(diss_t, key=lambda x: x[1])[:k] + diss_m_out = [id_ for _, id_ in diss_m_rst] + + source_ids.extend(ids) + id_m_out = [source_ids[i] for i, _ in diss_m_rst] + + return id_m_out, diss_m_out + def _do_merge(self, files_n_topk_results, topk, reverse=False, **kwargs): status = status_pb2.Status(error_code=status_pb2.SUCCESS, reason="Success") if not files_n_topk_results: return status, [] - request_results = defaultdict(list) + merge_id_results = [] + merge_dis_results = [] calc_time = time.time() for files_collection in files_n_topk_results: if isinstance(files_collection, tuple): status, _ = files_collection return status, [] - for request_pos, each_request_results in enumerate( - files_collection.topk_query_result): - request_results[request_pos].extend( - each_request_results.query_result_arrays) - request_results[request_pos] = sorted( - request_results[request_pos], - key=lambda x: x.distance, - reverse=reverse)[:topk] + + row_num = files_collection.row_num + ids = files_collection.ids + diss = files_collection.distances # distance collections + # TODO: batch_len is equal to topk, may need to compare with topk + batch_len = len(ids) // row_num + + for row_index in range(row_num): + id_batch = ids[row_index * batch_len: (row_index + 1) * batch_len] + dis_batch = diss[row_index * batch_len: (row_index + 1) * batch_len] + + if len(merge_id_results) < row_index: + raise ValueError("merge error") + elif len(merge_id_results) == row_index: + # TODO: may bug here + merge_id_results.append(id_batch) + merge_dis_results.append(dis_batch) + else: + merge_id_results[row_index], merge_dis_results[row_index] = \ + self._reduce(merge_id_results[row_index], id_batch, + merge_dis_results[row_index], dis_batch, + batch_len, + reverse) calc_time = time.time() - calc_time logger.info('Merge takes {}'.format(calc_time)) - results = sorted(request_results.items()) - topk_query_result = [] + id_mrege_list = [] + dis_mrege_list = [] - for result in results: - query_result = TopKQueryResult(query_result_arrays=result[1]) - topk_query_result.append(query_result) + for id_results, dis_results in zip(merge_id_results, merge_dis_results): + id_mrege_list.extend(id_results) + dis_mrege_list.extend(dis_results) - return status, topk_query_result + return status, id_mrege_list, dis_mrege_list def _do_query(self, context, @@ -109,8 +141,8 @@ class ServiceHandler(milvus_pb2_grpc.MilvusServiceServicer): file_ids=query_params['file_ids'], query_records=vectors, top_k=topk, - nprobe=nprobe, - lazy_=True) + nprobe=nprobe + ) end = time.time() logger.info('search_vectors_in_files takes: {}'.format(end - start)) @@ -241,7 +273,7 @@ class ServiceHandler(milvus_pb2_grpc.MilvusServiceServicer): logger.info('Search {}: topk={} nprobe={}'.format( table_name, topk, nprobe)) - metadata = {'resp_class': milvus_pb2.TopKQueryResultList} + metadata = {'resp_class': milvus_pb2.TopKQueryResult} if nprobe > self.MAX_NPROBE or nprobe <= 0: raise exceptions.InvalidArgumentError( @@ -275,22 +307,24 @@ class ServiceHandler(milvus_pb2_grpc.MilvusServiceServicer): query_range_array.append( Range(query_range.start_value, query_range.end_value)) - status, results = self._do_query(context, - table_name, - table_meta, - query_record_array, - topk, - nprobe, - query_range_array, - metadata=metadata) + status, id_results, dis_results = self._do_query(context, + table_name, + table_meta, + query_record_array, + topk, + nprobe, + query_range_array, + metadata=metadata) now = time.time() logger.info('SearchVector takes: {}'.format(now - start)) - topk_result_list = milvus_pb2.TopKQueryResultList( + topk_result_list = milvus_pb2.TopKQueryResult( status=status_pb2.Status(error_code=status.error_code, reason=status.reason), - topk_query_result=results) + row_num=len(query_record_array), + ids=id_results, + distances=dis_results) return topk_result_list @mark_grpc_method diff --git a/shards/requirements.txt b/shards/requirements.txt index 14bdde2a06..426ee0b704 100644 --- a/shards/requirements.txt +++ b/shards/requirements.txt @@ -14,8 +14,7 @@ py==1.8.0 pyasn1==0.4.7 pyasn1-modules==0.2.6 pylint==2.3.1 -pymilvus-test==0.2.28 -#pymilvus==0.2.0 +pymilvus==0.2.5 pyparsing==2.4.0 pytest==4.6.3 pytest-level==0.1.1