[feat][sdk] Impl vector get index metrics

DataCanvasIO · Mar 13, 2024 · cd9353c · cd9353c
1 parent 0265380
commit cd9353c
Show file tree

Hide file tree

Showing 14 changed files with 428 additions and 6 deletions.
diff --git a/src/example/sdk_vector_example.cc b/src/example/sdk_vector_example.cc
@@ -23,6 +23,7 @@
 #include "sdk/status.h"
 #include "sdk/vector.h"
 #include "sdk/vector/vector_common.h"
+#include "sdk/vector/vector_get_index_metrics_task.h"
 #include "sdk/vector/vector_index.h"
 #include "sdk/vector/vector_index_cache.h"
 
@@ -295,6 +296,25 @@ static void VectorScanQuery(bool use_index_name = false) {
   }
 }
 
+static void VectorGetIndexMetrics(bool use_index_name = false) {
+  Status tmp;
+  dingodb::sdk::IndexMetricsResult result;
+  if (use_index_name) {
+    tmp = g_vector_client->GetIndexMetricsByIndexName(g_schema_id, g_index_name, result);
+  } else {
+    tmp = g_vector_client->GetIndexMetricsByIndexId(g_index_id, result);
+  }
+
+  DINGO_LOG(INFO) << "vector get index metrics:" << tmp.ToString() << ", result :" << result.ToString();
+  if (tmp.ok()) {
+    CHECK_EQ(result.index_type, dingodb::sdk::VectorIndexType::kFlat);
+    CHECK_EQ(result.count, g_vector_ids.size());
+    CHECK_EQ(result.deleted_count, 0);
+    CHECK_EQ(result.max_vector_id, g_vector_ids[g_vector_ids.size() - 1]);
+    CHECK_EQ(result.min_vector_id, g_vector_ids[0]);
+  }
+}
+
 static void VectorDelete(bool use_index_name = false) {
   Status tmp;
   std::vector<dingodb::sdk::DeleteResult> result;
@@ -343,6 +363,7 @@ int main(int argc, char* argv[]) {
     VectorQuey();
     VectorGetBorder();
     VectorScanQuery();
+    VectorGetIndexMetrics();
     VectorDelete();
     VectorSearch();
 
@@ -358,6 +379,7 @@ int main(int argc, char* argv[]) {
     VectorQuey(true);
     VectorGetBorder(true);
     VectorScanQuery(true);
+    VectorGetIndexMetrics(true);
     VectorDelete(true);
     VectorSearch(true);
 

diff --git a/src/pysdk/dingosdk.swg b/src/pysdk/dingosdk.swg
@@ -84,6 +84,8 @@ namespace dingodb {
         %feature("docstring") VectorClient::GetBorderByIndexName "return Status, int64_t out_vector_id"
         %feature("docstring") VectorClient::ScanQueryByIndexId "return Status, ScanQueryResult out_result"
         %feature("docstring") VectorClient::ScanQueryByIndexName "return Status, ScanQueryResult out_result"
+        %feature("docstring") VectorClient::GetIndexMetricsByIndexId "return Status, IndexMetricsResult out_result"
+        %feature("docstring") VectorClient::GetIndexMetricsByIndexName "return Status, IndexMetricsResult out_result"
 
         %typemap(in, numinputs=0) Client** (Client* temp){
           temp = NULL;
@@ -162,15 +164,23 @@ namespace dingodb {
           $1 = &temp;
         }
         %typemap(argout) QueryResult& {
-          PyObject* obj = SWIG_NewPointerObj(new QueryResult(*$1), SWIGTYPE_p_dingodb__sdk__QueryResult, SWIG_POINTER_OWN);
+          PyObject* obj = SWIG_NewPointerObj(new QueryResult(*$1), $1_descriptor, SWIG_POINTER_OWN);
           $result = SWIG_AppendOutput($result, obj);
         }
 
         %typemap(in, numinputs=0) ScanQueryResult& (ScanQueryResult temp) {
           $1 = &temp;
         }
         %typemap(argout) ScanQueryResult& {
-          PyObject* obj = SWIG_NewPointerObj(new ScanQueryResult(*$1), SWIGTYPE_p_dingodb__sdk__ScanQueryResult, SWIG_POINTER_OWN);
+          PyObject* obj = SWIG_NewPointerObj(new ScanQueryResult(*$1), $1_descriptor, SWIG_POINTER_OWN);
+          $result = SWIG_AppendOutput($result, obj);
+        }
+
+        %typemap(in, numinputs=0) IndexMetricsResult& (IndexMetricsResult temp) {
+          $1 = &temp;
+        }
+        %typemap(argout) IndexMetricsResult& {
+          PyObject* obj = SWIG_NewPointerObj(new IndexMetricsResult(*$1), $1_descriptor, SWIG_POINTER_OWN);
           $result = SWIG_AppendOutput($result, obj);
         }
 

diff --git a/src/pysdk/pysdk_vector_example.py b/src/pysdk/pysdk_vector_example.py
@@ -181,6 +181,20 @@ def vector_scan_query(use_index_name=False):
         assert result.vectors[0].id == g_vector_ids[-1]
         assert result.vectors[1].id == g_vector_ids[-2]
 
+def vector_get_index_metrics(use_index_name=False):
+    if use_index_name:
+        tmp, result = g_vector_client.GetIndexMetricsByIndexName(g_schema_id, g_index_name)
+    else:
+        tmp, result = g_vector_client.GetIndexMetricsByIndexId(g_index_id)
+
+    print(f"vector get index metrics: {tmp.ToString()}, result : {result.ToString()}")
+    if tmp.ok():
+        assert result.index_type == dingosdk.kFlat
+        assert result.count == len(g_vector_ids)
+        assert result.deleted_count == 0
+        assert result.max_vector_id == g_vector_ids[-1]
+        assert result.min_vector_id == g_vector_ids[0]
+
 def vector_delete(use_index_name=False):
     if use_index_name:
         tmp, result = g_vector_client.DeleteByIndexName(g_schema_id, g_index_name, g_vector_ids)
@@ -202,6 +216,7 @@ def vector_delete(use_index_name=False):
     vector_query()
     vector_get_border()
     vector_scan_query()
+    vector_get_index_metrics()
     vector_delete()
     post_clean()
 
@@ -211,6 +226,7 @@ def vector_delete(use_index_name=False):
     vector_query(True)
     vector_get_border(True)
     vector_scan_query(True)
+    vector_get_index_metrics(True)
     vector_delete(True)
     post_clean(True)
 
diff --git a/src/sdk/CMakeLists.txt b/src/sdk/CMakeLists.txt
@@ -38,6 +38,7 @@ add_library(sdk
   vector/vector_batch_query_task.cc
   vector/vector_delete_task.cc
   vector/vector_get_border_task.cc
+  vector/vector_get_index_metrics_task.cc
   vector/vector_scan_query_task.cc
   vector/vector_search_task.cc
   utils/thread_pool_actuator.cc

diff --git a/src/sdk/vector.h b/src/sdk/vector.h
@@ -428,6 +428,17 @@ struct ScanQueryResult {
   std::string ToString() const;
 };
 
+struct IndexMetricsResult {
+  VectorIndexType index_type{kNoneIndexType};
+  int64_t count{0};
+  int64_t deleted_count{0};
+  int64_t max_vector_id{0};
+  int64_t min_vector_id{0};
+  int64_t memory_bytes{0};
+
+  std::string ToString() const;
+};
+
 class VectorIndexCreator {
  public:
   ~VectorIndexCreator();
@@ -496,6 +507,9 @@ class VectorClient {
   Status ScanQueryByIndexName(int64_t schema_id, const std::string& index_name, const ScanQueryParam& query_param,
                               ScanQueryResult& out_result);
 
+  Status GetIndexMetricsByIndexId(int64_t index_id, IndexMetricsResult& out_result);
+  Status GetIndexMetricsByIndexName(int64_t schema_id, const std::string& index_name, IndexMetricsResult& out_result);
+
  private:
   friend class Client;
 

diff --git a/src/sdk/vector/index_service_rpc.cc b/src/sdk/vector/index_service_rpc.cc
@@ -36,6 +36,7 @@ DEFINE_INDEX_SERVICE_RPC(VectorDelete);
 DEFINE_INDEX_SERVICE_RPC(VectorBatchQuery);
 DEFINE_INDEX_SERVICE_RPC(VectorGetBorderId);
 DEFINE_INDEX_SERVICE_RPC(VectorScanQuery);
+DEFINE_INDEX_SERVICE_RPC(VectorGetRegionMetrics);
 
 }  // namespace sdk
 }  // namespace dingodb
diff --git a/src/sdk/vector/index_service_rpc.h b/src/sdk/vector/index_service_rpc.h
@@ -45,6 +45,7 @@ DECLARE_INDEX_SERVICE_RPC(VectorDelete);
 DECLARE_INDEX_SERVICE_RPC(VectorBatchQuery);
 DECLARE_INDEX_SERVICE_RPC(VectorGetBorderId);
 DECLARE_INDEX_SERVICE_RPC(VectorScanQuery);
+DECLARE_INDEX_SERVICE_RPC(VectorGetRegionMetrics);
 
 }  // namespace sdk
 }  // namespace dingodb

diff --git a/src/sdk/vector/vector_client.cc b/src/sdk/vector/vector_client.cc
@@ -21,6 +21,7 @@
 #include "sdk/vector/vector_batch_query_task.h"
 #include "sdk/vector/vector_delete_task.h"
 #include "sdk/vector/vector_get_border_task.h"
+#include "sdk/vector/vector_get_index_metrics_task.h"
 #include "sdk/vector/vector_index_cache.h"
 #include "sdk/vector/vector_scan_query_task.h"
 #include "sdk/vector/vector_search_task.h"
@@ -125,6 +126,21 @@ Status VectorClient::ScanQueryByIndexName(int64_t schema_id, const std::string&
   VectorScanQueryTask task(stub_, index_id, query_param, out_result);
   return task.Run();
 }
+
+Status VectorClient::GetIndexMetricsByIndexId(int64_t index_id, IndexMetricsResult& out_result) {
+  VectorGetIndexMetricsTask task(stub_, index_id, out_result);
+  return task.Run();
+}
+
+Status VectorClient::GetIndexMetricsByIndexName(int64_t schema_id, const std::string& index_name,
+                                                IndexMetricsResult& out_result) {
+  int64_t index_id{0};
+  DINGO_RETURN_NOT_OK(
+      stub_.GetVectorIndexCache()->GetIndexIdByKey(EncodeVectorIndexCacheKey(schema_id, index_name), index_id));
+  CHECK_GT(index_id, 0);
+  VectorGetIndexMetricsTask task(stub_, index_id, out_result);
+  return task.Run();
+}
 }  // namespace sdk
 
 }  // namespace dingodb
diff --git a/src/sdk/vector/vector_common.h b/src/sdk/vector/vector_common.h
@@ -272,6 +272,18 @@ static VectorWithDistance InternalVectorWithDistance2VectorWithDistance(const pb
   return std::move(to_return);
 }
 
+static IndexMetricsResult InternalVectorIndexMetrics2IndexMetricsResult(const pb::common::VectorIndexMetrics& pb) {
+  IndexMetricsResult to_return;
+  to_return.index_type = InternalVectorIndexTypePB2VectorIndexType(pb.vector_index_type());
+  to_return.count = pb.current_count();
+  to_return.deleted_count = pb.deleted_count();
+  to_return.max_vector_id = pb.max_id();
+  to_return.min_vector_id = pb.min_id();
+  to_return.memory_bytes = pb.memory_bytes();
+
+  return to_return;
+}
+
 static void FillSearchFlatParamPB(pb::common::SearchFlatParam* pb, const SearchParam& parameter) {
   if (parameter.extra_params.find(SearchExtraParamType::kParallelOnQueries) != parameter.extra_params.end()) {
     pb->set_parallel_on_queries(parameter.extra_params.at(SearchExtraParamType::kParallelOnQueries));

diff --git a/src/sdk/vector/vector_get_border_task.h b/src/sdk/vector/vector_get_border_task.h
@@ -30,7 +30,7 @@ class VectorGetBorderPartTask;
 class VectorGetBorderTask : public VectorTask {
  public:
   VectorGetBorderTask(const ClientStub& stub, int64_t index_id, bool is_max, int64_t& out_vector_id)
-      : VectorTask(stub), index_id_(index_id), is_max_(is_max), out_vector_id_(out_vector_id), controller_(stub, rpc_) {
+      : VectorTask(stub), index_id_(index_id), is_max_(is_max), out_vector_id_(out_vector_id) {
     target_vector_id_ = is_max_ ? -1 : INT64_MAX;
   }
 
@@ -50,9 +50,6 @@ class VectorGetBorderTask : public VectorTask {
   std::shared_ptr<VectorIndex> vector_index_;
   int64_t target_vector_id_;
 
-  VectorGetBorderIdRpc rpc_;
-  StoreRpcController controller_;
-
   std::shared_mutex rw_lock_;
   std::set<int64_t> next_part_ids_;
   Status status_;