[feat][sdk] Impl vector count

DataCanvasIO · Mar 13, 2024 · 79f60a8 · 79f60a8
1 parent 3d65704
commit 79f60a8
Show file tree

Hide file tree

Showing 10 changed files with 420 additions and 2 deletions.
diff --git a/src/example/sdk_vector_example.cc b/src/example/sdk_vector_example.cc
@@ -315,6 +315,57 @@ static void VectorGetIndexMetrics(bool use_index_name = false) {
   }
 }
 
+static void VectorCount(bool use_index_name = false) {
+  {
+    Status tmp;
+    int64_t result{0};
+    if (use_index_name) {
+      tmp = g_vector_client->CountByIndexName(g_schema_id, g_index_name, 0, g_vector_ids[g_vector_ids.size() - 1] + 1,
+                                              result);
+    } else {
+      tmp = g_vector_client->CountByIndexId(g_index_id, 0, g_vector_ids[g_vector_ids.size() - 1] + 1, result);
+    }
+
+    DINGO_LOG(INFO) << "vector count:" << tmp.ToString() << ", result :" << result;
+    if (tmp.ok()) {
+      CHECK_EQ(result, g_vector_ids.size());
+    }
+  }
+
+  {
+    Status tmp;
+    int64_t result{0};
+    int64_t start_vector_id = g_vector_ids[g_vector_ids.size() - 1] + 1;
+    int64_t end_vector_id = start_vector_id + 1;
+    if (use_index_name) {
+      tmp = g_vector_client->CountByIndexName(g_schema_id, g_index_name, start_vector_id, end_vector_id, result);
+    } else {
+      tmp = g_vector_client->CountByIndexId(g_index_id, start_vector_id, end_vector_id, result);
+    }
+
+    DINGO_LOG(INFO) << "vector count:" << tmp.ToString() << ", result :" << result;
+    if (tmp.ok()) {
+      CHECK_EQ(result, 0);
+    }
+  }
+
+  {
+    Status tmp;
+    int64_t result{0};
+    if (use_index_name) {
+      tmp = g_vector_client->CountByIndexName(g_schema_id, g_index_name, g_vector_ids[0],
+                                              g_vector_ids[g_vector_ids.size() - 1], result);
+    } else {
+      tmp = g_vector_client->CountByIndexId(g_index_id, g_vector_ids[0], g_vector_ids[g_vector_ids.size() - 1], result);
+    }
+
+    DINGO_LOG(INFO) << "vector count:" << tmp.ToString() << ", result :" << result;
+    if (tmp.ok()) {
+      CHECK_EQ(result, g_vector_ids.size() - 1);
+    }
+  }
+}
+
 static void VectorDelete(bool use_index_name = false) {
   Status tmp;
   std::vector<dingodb::sdk::DeleteResult> result;
@@ -364,6 +415,7 @@ int main(int argc, char* argv[]) {
     VectorGetBorder();
     VectorScanQuery();
     VectorGetIndexMetrics();
+    VectorCount();
     VectorDelete();
     VectorSearch();
 
@@ -380,6 +432,7 @@ int main(int argc, char* argv[]) {
     VectorGetBorder(true);
     VectorScanQuery(true);
     VectorGetIndexMetrics(true);
+    VectorCount(true);
     VectorDelete(true);
     VectorSearch(true);
 

diff --git a/src/pysdk/dingosdk.swg b/src/pysdk/dingosdk.swg
@@ -43,6 +43,7 @@ using namespace dingodb::sdk;
 %apply bool& OUTPUT { bool& out_create_in_progress };
 %apply int64_t& OUTPUT  { int64_t& out_index_id };
 %apply int64_t& OUTPUT  { int64_t& out_vector_id };
+%apply int64_t& OUTPUT  { int64_t& out_count };
 
 // rawkv
 %apply std::string& OUTPUT { std::string& out_value };
@@ -86,6 +87,8 @@ namespace dingodb {
         %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"
+        %feature("docstring") VectorClient::CountByIndexId "return Status, int64_t out_count"
+        %feature("docstring") VectorClient::CountByIndexName "return Status, int64_t out_count"
 
         %typemap(in, numinputs=0) Client** (Client* temp){
           temp = NULL;
@@ -142,7 +145,7 @@ namespace dingodb {
         %typemap(argout) std::vector<SearchResult>& {
             PyObject* obj = PyList_New($1->size());
             for (size_t i = 0; i < $1->size(); i++) {
-                PyObject* item = SWIG_NewPointerObj(new SearchResult((*$1)[i]), SWIGTYPE_p_dingodb__sdk__SearchResult, SWIG_POINTER_OWN);
+                PyObject* item = SWIG_NewPointerObj(new SearchResult((*$1)[i]), $descriptor(dingodb::sdk::SearchResult*), SWIG_POINTER_OWN);
                 PyList_SetItem(obj, i, item);
             }
             $result = SWIG_AppendOutput($result, obj);
@@ -154,7 +157,7 @@ namespace dingodb {
         %typemap(argout) std::vector<DeleteResult>& {
           PyObject* obj = PyList_New($1->size());
           for (size_t i = 0; i < $1->size(); i++) {
-            PyObject* item = SWIG_NewPointerObj(new DeleteResult((*$1)[i]), SWIGTYPE_p_dingodb__sdk__DeleteResult, SWIG_POINTER_OWN);
+            PyObject* item = SWIG_NewPointerObj(new DeleteResult((*$1)[i]), $descriptor(dingodb::sdk::DeleteResult*), SWIG_POINTER_OWN);
             PyList_SetItem(obj, i, item);
           }
           $result = SWIG_AppendOutput($result, obj);

diff --git a/src/pysdk/pysdk_vector_example.py b/src/pysdk/pysdk_vector_example.py
@@ -195,6 +195,39 @@ def vector_get_index_metrics(use_index_name=False):
         assert result.max_vector_id == g_vector_ids[-1]
         assert result.min_vector_id == g_vector_ids[0]
 
+def vector_count(use_index_name=False):
+    result = 0
+    if use_index_name:
+        tmp, result = g_vector_client.CountByIndexName(g_schema_id, g_index_name, 0, g_vector_ids[-1] + 1)
+    else:
+        tmp, result= g_vector_client.CountByIndexId(g_index_id, 0, g_vector_ids[-1] + 1)
+
+    print(f"vector count: {tmp.ToString()}, result: {result}")
+    if tmp.ok():
+        assert result == len(g_vector_ids)
+
+    result = 0
+    start_vector_id = g_vector_ids[-1] + 1
+    end_vector_id = start_vector_id + 1
+    if use_index_name:
+        tmp, result = g_vector_client.CountByIndexName(g_schema_id, g_index_name, start_vector_id, end_vector_id)
+    else:
+        tmp, result= g_vector_client.CountByIndexId(g_index_id, start_vector_id, end_vector_id)
+
+    print(f"vector count: {tmp.ToString()}, result: {result}")
+    if tmp.ok():
+        assert result == 0
+
+    result = 0
+    if use_index_name:
+        tmp, result = g_vector_client.CountByIndexName(g_schema_id, g_index_name, g_vector_ids[0], g_vector_ids[-1])
+    else:
+        tmp, result = g_vector_client.CountByIndexId(g_index_id, g_vector_ids[0], g_vector_ids[-1])
+
+    print(f"vector count: {tmp.ToString()}, result: {result}")
+    if tmp.ok():
+        assert result == len(g_vector_ids) - 1
+
 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)
@@ -217,6 +250,7 @@ def vector_delete(use_index_name=False):
     vector_get_border()
     vector_scan_query()
     vector_get_index_metrics()
+    vector_count()
     vector_delete()
     post_clean()
 
@@ -227,6 +261,7 @@ def vector_delete(use_index_name=False):
     vector_get_border(True)
     vector_scan_query(True)
     vector_get_index_metrics(True)
+    vector_count(True)
     vector_delete(True)
     post_clean(True)
 
diff --git a/src/sdk/CMakeLists.txt b/src/sdk/CMakeLists.txt
@@ -36,6 +36,7 @@ add_library(sdk
   vector/vector_task.cc
   vector/vector_add_task.cc
   vector/vector_batch_query_task.cc
+  vector/vector_count_task.cc
   vector/vector_delete_task.cc
   vector/vector_get_border_task.cc
   vector/vector_get_index_metrics_task.cc

diff --git a/src/sdk/vector.h b/src/sdk/vector.h
@@ -510,6 +510,9 @@ class VectorClient {
   Status GetIndexMetricsByIndexId(int64_t index_id, IndexMetricsResult& out_result);
   Status GetIndexMetricsByIndexName(int64_t schema_id, const std::string& index_name, IndexMetricsResult& out_result);
 
+  Status CountByIndexId(int64_t index_id, int64_t start_vector_id, int64_t end_vector_id, int64_t& out_count);
+  Status CountByIndexName(int64_t schema_id, const std::string& index_name, int64_t start_vector_id, int64_t end_vector_id, int64_t& out_count);
+
  private:
   friend class Client;
 

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

diff --git a/src/sdk/vector/vector_client.cc b/src/sdk/vector/vector_client.cc
@@ -19,6 +19,7 @@
 #include "sdk/vector.h"
 #include "sdk/vector/vector_add_task.h"
 #include "sdk/vector/vector_batch_query_task.h"
+#include "sdk/vector/vector_count_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"
@@ -141,6 +142,24 @@ Status VectorClient::GetIndexMetricsByIndexName(int64_t schema_id, const std::st
   VectorGetIndexMetricsTask task(stub_, index_id, out_result);
   return task.Run();
 }
+
+Status VectorClient::CountByIndexId(int64_t index_id, int64_t start_vector_id, int64_t end_vector_id,
+                                    int64_t& out_count) {
+  VectorCountTask task(stub_, index_id, start_vector_id, end_vector_id, out_count);
+  return task.Run();
+}
+
+Status VectorClient::CountByIndexName(int64_t schema_id, const std::string& index_name, int64_t start_vector_id,
+                                      int64_t end_vector_id, int64_t& out_count) {
+  int64_t index_id{0};
+  DINGO_RETURN_NOT_OK(
+      stub_.GetVectorIndexCache()->GetIndexIdByKey(EncodeVectorIndexCacheKey(schema_id, index_name), index_id));
+  CHECK_GT(index_id, 0);
+
+  VectorCountTask task(stub_, index_id, start_vector_id, end_vector_id, out_count);
+  return task.Run();
+}
+
 }  // namespace sdk
 
 }  // namespace dingodb