From 17fedaf6985ead10b0256c796bc5310d6b7eba16 Mon Sep 17 00:00:00 2001 From: yukun Date: Mon, 24 Aug 2020 23:34:15 +0800 Subject: [PATCH] Add partition_tag in GetPageEntities (#3434) * Fix TestSearchDSL multi fields bug Signed-off-by: fishpenguin * Fix set_config bug Signed-off-by: fishpenguin * Add partition_tag in GetPageEntities Signed-off-by: fishpenguin --- core/src/server/web_impl/dto/VectorDto.hpp | 4 ++++ .../web_impl/handler/WebRequestHandler.cpp | 23 ++++++++++++++----- .../web_impl/handler/WebRequestHandler.h | 4 ++-- core/unittest/server/test_web.cpp | 5 ++-- 4 files changed, 26 insertions(+), 10 deletions(-) diff --git a/core/src/server/web_impl/dto/VectorDto.hpp b/core/src/server/web_impl/dto/VectorDto.hpp index 8adf2fe06..29d5b1eb2 100644 --- a/core/src/server/web_impl/dto/VectorDto.hpp +++ b/core/src/server/web_impl/dto/VectorDto.hpp @@ -22,6 +22,10 @@ class EntityIdsDto : public ODTO { DTO_INIT(EntityIdsDto, DTO) DTO_FIELD(List, ids); + + DTO_FIELD(String, message); + + DTO_FIELD(Int64, code); }; #include OATPP_CODEGEN_END(DTO) diff --git a/core/src/server/web_impl/handler/WebRequestHandler.cpp b/core/src/server/web_impl/handler/WebRequestHandler.cpp index b9ccab832..709e787a4 100644 --- a/core/src/server/web_impl/handler/WebRequestHandler.cpp +++ b/core/src/server/web_impl/handler/WebRequestHandler.cpp @@ -307,8 +307,8 @@ WebRequestHandler::GetCollectionStat(const std::string& collection_name, nlohman } Status -WebRequestHandler::GetPageEntities(const std::string& collection_name, const int64_t page_size, const int64_t offset, - nlohmann::json& json_out) { +WebRequestHandler::GetPageEntities(const std::string& collection_name, const std::string& partition_tag, + const int64_t page_size, const int64_t offset, nlohmann::json& json_out) { std::string collection_info; STATUS_CHECK(req_handler_.GetCollectionStats(context_ptr_, collection_name, collection_info)); nlohmann::json json_info = nlohmann::json::parse(collection_info); @@ -324,6 +324,9 @@ WebRequestHandler::GetPageEntities(const std::string& collection_name, const int int64_t row_count = 0; bool already_find = false; for (auto& json_partition : json_info["partitions"]) { + if (!partition_tag.empty() && json_partition["tag"] != partition_tag) { + continue; + } for (auto& json_segment : json_partition["segments"]) { auto count = json_segment["row_count"].get(); row_count += count; @@ -333,7 +336,9 @@ WebRequestHandler::GetPageEntities(const std::string& collection_name, const int if (not already_find && entity_num >= count) { entity_num -= count; } - segment_ids.emplace_back(json_segment["id"].get()); + if (already_find) { + segment_ids.emplace_back(json_segment["id"].get()); + } } } json_out["count"] = row_count; @@ -1731,6 +1736,8 @@ WebRequestHandler::InsertEntity(const OString& collection_name, const milvus::se ids_dto->ids->push_back(std::to_string(pdata[i]).c_str()); } } + ids_dto->code = status.code(); + ids_dto->message = status.message().c_str(); ASSIGN_RETURN_STATUS_DTO(status) } @@ -1745,7 +1752,11 @@ WebRequestHandler::GetEntity(const milvus::server::web::OString& collection_name nlohmann::json json_out; auto offset = std::stoi(query_params.get("offset")->std_str(), nullptr); auto page_size = std::stoi(query_params.get("page_size")->std_str(), nullptr); - status = GetPageEntities(collection_name->std_str(), page_size, offset, json_out); + std::string partition_tag; + if (query_params.get("partition_tag")) { + partition_tag = query_params.get("partition_tag")->std_str(); + } + status = GetPageEntities(collection_name->std_str(), partition_tag, page_size, offset, json_out); AddStatusToJson(json_out, status.code(), status.message()); response = json_out.dump().c_str(); ASSIGN_RETURN_STATUS_DTO(status); @@ -1780,9 +1791,9 @@ WebRequestHandler::GetEntity(const milvus::server::web::OString& collection_name nlohmann::json json; AddStatusToJson(json, status.code(), status.message()); if (entity_result_json.empty()) { - json["entities"] = std::vector(); + json = std::vector(); } else { - json["entities"] = entity_result_json; + json = entity_result_json; } response = json.dump().c_str(); } catch (std::exception& e) { diff --git a/core/src/server/web_impl/handler/WebRequestHandler.h b/core/src/server/web_impl/handler/WebRequestHandler.h index 082ba58c6..1bdab72e4 100644 --- a/core/src/server/web_impl/handler/WebRequestHandler.h +++ b/core/src/server/web_impl/handler/WebRequestHandler.h @@ -103,8 +103,8 @@ class WebRequestHandler { nlohmann::json& json_out); Status - GetPageEntities(const std::string& collection_name, const int64_t page_size, const int64_t offset, - nlohmann::json& json_out); + GetPageEntities(const std::string& collection_name, const std::string& partition_tag, const int64_t page_size, + const int64_t offset, nlohmann::json& json_out); Status GetSegmentIds(const std::string& collection_name, int64_t segment_id, int64_t page_size, int64_t offset, diff --git a/core/unittest/server/test_web.cpp b/core/unittest/server/test_web.cpp index 72221228b..396fbc368 100644 --- a/core/unittest/server/test_web.cpp +++ b/core/unittest/server/test_web.cpp @@ -288,7 +288,8 @@ class TestClient : public oatpp::web::client::ApiClient { PATH(String, collection_name, "collection_name")) API_CALL("GET", "/collections/{collection_name}/entities", getEntity, - PATH(String, collection_name, "collection_name"), QUERY(String, offset), QUERY(String, page_size)) + PATH(String, collection_name, "collection_name"), QUERY(String, offset), QUERY(String, page_size), + QUERY(String, partition_tag)) API_CALL("GET", "/collections/{collection_name}/entities", getEntityByID, PATH(String, collection_name, "collection_name"), QUERY(String, ids)) @@ -717,7 +718,7 @@ TEST_F(WebControllerTest, GET_PAGE_ENTITY) { std::string offset = "0"; std::string page_size = "10"; - response = client_ptr->getEntity(collection_name.c_str(), offset.c_str(), page_size.c_str(), connection_ptr); + response = client_ptr->getEntity(collection_name.c_str(), offset.c_str(), page_size.c_str(), "", connection_ptr); ASSERT_EQ(OStatus::CODE_200.code, response->getStatusCode()); // offset = "10"; -- GitLab