get collection info (#2991)

Signed-off-by: Nyhmo <yihua.mo@zilliz.com>

core/src/db/SSDB.h

@@ -62,6 +62,9 @@ class SSDB {
     virtual Status
     AllCollections(std::vector<std::string>& names) = 0;
 
+    virtual Status
+    GetCollectionInfo(const std::string& collection_name, std::string& collection_info) = 0;
+
     virtual Status
     GetCollectionRowCount(const std::string& collection_name, uint64_t& row_count) = 0;
 

core/src/db/SSDBImpl.cpp

@@ -283,6 +283,20 @@ SSDBImpl::AllCollections(std::vector<std::string>& names) {
     return snapshot::Snapshots::GetInstance().GetCollectionNames(names);
 }
 
+Status
+SSDBImpl::GetCollectionInfo(const std::string& collection_name, std::string& collection_info) {
+    CHECK_INITIALIZED;
+
+    nlohmann::json json;
+    auto status = GetSnapshotInfo(collection_name, json);
+    if (!status.ok()) {
+        return status;
+    }
+
+    collection_info = json.dump();
+    return Status::OK();
+}
+
 Status
 SSDBImpl::GetCollectionRowCount(const std::string& collection_name, uint64_t& row_count) {
     CHECK_INITIALIZED;

core/src/db/SSDBImpl.h

@@ -57,6 +57,9 @@ class SSDBImpl : public SSDB {
     Status
     AllCollections(std::vector<std::string>& names) override;
 
+    Status
+    GetCollectionInfo(const std::string& collection_name, std::string& collection_info);
+
     Status
     GetCollectionRowCount(const std::string& collection_name, uint64_t& row_count) override;
 

core/src/db/SnapshotUtils.cpp

@@ -10,18 +10,33 @@
 // or implied. See the License for the specific language governing permissions and limitations under the License.
 
 #include "db/SnapshotUtils.h"
+#include "db/SnapshotHandlers.h"
+#include "db/SnapshotVisitor.h"
 #include "db/snapshot/CompoundOperations.h"
 #include "db/snapshot/Resources.h"
 #include "db/snapshot/Snapshots.h"
 #include "segment/Segment.h"
 
 #include <memory>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
 namespace milvus {
 namespace engine {
 
+namespace {
+constexpr const char* JSON_ROW_COUNT = "row_count";
+constexpr const char* JSON_ID = "id";
+constexpr const char* JSON_PARTITIONS = "partitions";
+constexpr const char* JSON_PARTITION_TAG = "tag";
+constexpr const char* JSON_SEGMENTS = "segments";
+constexpr const char* JSON_NAME = "name";
+constexpr const char* JSON_FIELDS = "fields";
+constexpr const char* JSON_INDEX_NAME = "index_name";
+constexpr const char* JSON_DATA_SIZE = "data_size";
+}  // namespace
+
 Status
 SetSnapshotIndex(const std::string& collection_name, const std::string& field_name,
                  engine::CollectionIndex& index_info) {
@@ -128,5 +143,103 @@ IsVectorField(const engine::snapshot::FieldPtr& field) {
     return ftype == engine::FIELD_TYPE::VECTOR_FLOAT || ftype == engine::FIELD_TYPE::VECTOR_BINARY;
 }
 
+Status
+GetSnapshotInfo(const std::string& collection_name, nlohmann::json& json_info) {
+    snapshot::ScopedSnapshotT ss;
+    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));
+
+    size_t total_row_count = 0;
+
+    std::unordered_map<snapshot::ID_TYPE, milvus::json> partitions;
+    auto partition_names = ss->GetPartitionNames();
+    for (auto& name : partition_names) {
+        auto partition = ss->GetPartition(name);
+
+        milvus::json json_partition;
+        json_partition[JSON_PARTITION_TAG] = name;
+        json_partition[JSON_ID] = partition->GetID();
+
+        auto partition_commit = ss->GetPartitionCommitByPartitionId(partition->GetID());
+        json_partition[JSON_ROW_COUNT] = partition_commit->GetRowCount();
+        total_row_count += partition_commit->GetRowCount();
+
+        partitions.insert(std::make_pair(partition->GetID(), json_partition));
+    }
+
+    snapshot::IDS_TYPE segment_ids;
+    auto handler = std::make_shared<SegmentsToSearchCollector>(ss, segment_ids);
+    handler->Iterate();
+
+    std::unordered_map<snapshot::ID_TYPE, std::vector<milvus::json>> segments;
+    for (auto id : segment_ids) {
+        auto segment_commit = ss->GetSegmentCommitBySegmentId(id);
+        if (segment_commit == nullptr) {
+            continue;
+        }
+
+        milvus::json json_fields;
+        auto seg_visitor = engine::SegmentVisitor::Build(ss, id);
+        auto& field_visitors = seg_visitor->GetFieldVisitors();
+        for (auto& iter : field_visitors) {
+            milvus::json json_field;
+            const engine::snapshot::FieldPtr& field = iter.second->GetField();
+            json_field[JSON_NAME] = field->GetName();
+
+            std::string index_name;
+            uint64_t total_size = 0;
+            auto element_visitor = iter.second->GetElementVisitor(engine::FieldElementType::FET_RAW);
+            if (element_visitor) {
+                if (element_visitor->GetFile()) {
+                    total_size += element_visitor->GetFile()->GetSize();
+                }
+                if (element_visitor->GetElement()) {
+                    index_name = element_visitor->GetElement()->GetName();
+                }
+            }
+
+            auto index_visitor = iter.second->GetElementVisitor(engine::FieldElementType::FET_INDEX);
+            if (index_visitor && index_visitor->GetFile()) {
+                if (index_visitor->GetFile()) {
+                    total_size += index_visitor->GetFile()->GetSize();
+                }
+                if (index_visitor->GetElement()) {
+                    index_name = index_visitor->GetElement()->GetName();
+                }
+            }
+
+            auto compress_visitor = iter.second->GetElementVisitor(engine::FieldElementType::FET_COMPRESS_SQ8);
+            if (compress_visitor && compress_visitor->GetFile()) {
+                total_size += compress_visitor->GetFile()->GetSize();
+            }
+
+            json_field[JSON_INDEX_NAME] = index_name;
+            json_field[JSON_DATA_SIZE] = total_size;
+        }
+
+        milvus::json json_segment;
+        json_segment[JSON_ID] = id;
+        json_segment[JSON_ROW_COUNT] = segment_commit->GetRowCount();
+        json_segment[JSON_DATA_SIZE] = segment_commit->GetSize();
+        json_segment[JSON_FIELDS] = json_fields;
+        segments[segment_commit->GetPartitionId()].push_back(json_segment);
+    }
+
+    milvus::json json_partitions;
+    for (auto pair : partitions) {
+        milvus::json json_segments;
+        auto seg_array = segments[pair.first];
+        for (auto& json : seg_array) {
+            json_segments.push_back(json);
+        }
+        pair.second[JSON_SEGMENTS] = json_segments;
+        json_partitions.push_back(pair.second);
+    }
+
+    json_info[JSON_ROW_COUNT] = total_row_count;
+    json_info[JSON_PARTITIONS] = json_partitions;
+
+    return Status::OK();
+}
+
 }  // namespace engine
 }  // namespace milvus

core/src/db/SnapshotUtils.h

@@ -13,6 +13,7 @@
 
 #include "db/Types.h"
 #include "db/snapshot/Resources.h"
+#include "thirdparty/nlohmann/json.hpp"
 
 #include <string>
 
@@ -33,5 +34,8 @@ DeleteSnapshotIndex(const std::string& collection_name, const std::string& field
 bool
 IsVectorField(const engine::snapshot::FieldPtr& field);
 
+Status
+GetSnapshotInfo(const std::string& collection_name, nlohmann::json& json_info);
+
 }  // namespace engine
 }  // namespace milvus

core/src/server/delivery/request/ShowCollectionInfoRequest.cpp

@@ -66,8 +66,7 @@ ShowCollectionInfoRequest::OnExecute() {
     }
 
     // step 3: get partitions
-    // TODO(yukun): SSDBImpl::GetCollectionInfo has not implemented yet
-    status = DBWrapper::DB()->GetCollectionInfo(collection_name_, collection_info_);
+    status = DBWrapper::SSDB()->GetCollectionInfo(collection_name_, collection_info_);
     if (!status.ok()) {
         return status;
     }