[feat][sdk] Support create vector index with auto incre

yuhaijun999 · Apr 10, 2024 · 18171f9 · 18171f9
1 parent 3007af7
commit 18171f9
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 8 deletions.
diff --git a/src/example/sdk_vector_example.cc b/src/example/sdk_vector_example.cc
@@ -54,6 +54,8 @@ static void PrepareVectorIndex() {
                       .SetReplicaNum(3)
                       .SetRangePartitions(g_range_partition_seperator_ids)
                       .SetFlatParam(g_flat_param)
+                      .SetAutoIncrement(true)
+                      .SetAutoIncrementStart(1)
                       .Create(g_index_id);
   DINGO_LOG(INFO) << "Create index status: " << create.ToString() << ", index_id:" << g_index_id;
   sleep(20);

diff --git a/src/sdk/vector.h b/src/sdk/vector.h
@@ -128,6 +128,7 @@ struct Vector {
   int32_t dimension;
   ValueType value_type;
   std::vector<float> float_values;
+  // TODO: support
   std::vector<uint8_t> binary_values;
 
   explicit Vector() : value_type(kNoneValueType), dimension(0) {}
@@ -462,9 +463,9 @@ class VectorIndexCreator {
   // VectorIndexCreator& SetDiskAnnParam(DiskAnnParam& params);
   VectorIndexCreator& SetBruteForceParam(const BruteForceParam& params);
 
-  // VectorIndexCreator& SetAutoIncrement(bool auto_incr);
-
-  // VectorIndexCreator& SetAutoIncrementStart(int64_t start_id);
+  VectorIndexCreator& SetAutoIncrement(bool auto_incr);
+  // start_id should greater than 0, when set auto_increment is set to true
+  VectorIndexCreator& SetAutoIncrementStart(int64_t start_id);
 
   Status Create(int64_t& out_index_id);
 

diff --git a/src/sdk/vector/vector_index_creator.cc b/src/sdk/vector/vector_index_creator.cc
@@ -79,6 +79,17 @@ VectorIndexCreator& VectorIndexCreator::SetBruteForceParam(const BruteForceParam
   return *this;
 }
 
+VectorIndexCreator& VectorIndexCreator::SetAutoIncrement(bool auto_incr) {
+  data_->auto_incr = auto_incr;
+  return *this;
+}
+
+VectorIndexCreator& VectorIndexCreator::SetAutoIncrementStart(int64_t start_id) {
+  data_->auto_incr = true;
+  data_->auto_incr_start = start_id;
+  return *this;
+}
+
 // TODO: check partition is illegal
 // TODO: support hash partitions
 Status VectorIndexCreator::Create(int64_t& out_index_id) {
@@ -123,8 +134,14 @@ Status VectorIndexCreator::Create(int64_t& out_index_id) {
   auto* index_definition_pb = request.mutable_index_definition();
   index_definition_pb->set_name(data_->index_name);
   index_definition_pb->set_replica(data_->replica_num);
-  // index_definition->set_with_auto_incrment(true);
-  // index_definition->set_auto_increment(1024);
+  index_definition_pb->set_with_auto_incrment(data_->auto_incr);
+  if (data_->auto_incr && data_->auto_incr_start.has_value()) {
+    int64_t start = data_->auto_incr_start.value();
+    if (start <= 0) {
+      return Status::InvalidArgument("auto_increment_start must greater 0");
+    }
+    index_definition_pb->set_auto_increment(start);
+  }
 
   // vector index parameter
   data_->BuildIndexParameter(index_definition_pb->mutable_index_parameter());

diff --git a/src/sdk/vector/vector_index_creator_internal_data.h b/src/sdk/vector/vector_index_creator_internal_data.h
@@ -37,7 +37,6 @@ class VectorIndexCreator::Data {
         version(1),
         replica_num(3),
         index_type(kNoneIndexType),
-        auto_incr(false),
         wait(true) {}
 
   ~Data() = default;
@@ -86,9 +85,9 @@ class VectorIndexCreator::Data {
   std::optional<DiskAnnParam> diskann_param;
   std::optional<BruteForceParam> brute_force_param;
 
-  // TODO: Support
-  bool auto_incr;
+  bool auto_incr{false};
   std::optional<int64_t> auto_incr_start;
+
   bool wait;
 };