[feat][sdk] Support vector scan query with filter

DataCanvasIO · Apr 9, 2024 · 39d0afe · 39d0afe
1 parent 7428045
commit 39d0afe
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 8 deletions.
diff --git a/src/example/sdk_vector_example.cc b/src/example/sdk_vector_example.cc
@@ -296,7 +296,7 @@ static void VectorQuey(bool use_index_name = false) {
     query = g_vector_client->BatchQueryByIndexId(g_index_id, param, result);
   }
 
-  DINGO_LOG(INFO) << "vector query:" << query.ToString();
+  DINGO_LOG(INFO) << "vector query: " << query.ToString();
   DINGO_LOG(INFO) << "vector query result:" << result.ToString();
   CHECK_EQ(result.vectors.size(), g_vector_ids.size());
 }
@@ -351,7 +351,7 @@ static void VectorScanQuery(bool use_index_name = false) {
       tmp = g_vector_client->ScanQueryByIndexId(g_index_id, param, result);
     }
 
-    DINGO_LOG(INFO) << "vector scan query:" << tmp.ToString() << ", result:" << result.ToString();
+    DINGO_LOG(INFO) << "vector forward scan query: " << tmp.ToString() << ", result:" << result.ToString();
     if (tmp.ok()) {
       CHECK_EQ(result.vectors[0].id, g_vector_ids[0]);
       CHECK_EQ(result.vectors[1].id, g_vector_ids[1]);
@@ -374,12 +374,45 @@ static void VectorScanQuery(bool use_index_name = false) {
       tmp = g_vector_client->ScanQueryByIndexId(g_index_id, param, result);
     }
 
-    DINGO_LOG(INFO) << "vector scan query:" << tmp.ToString() << ", result:" << result.ToString();
+    DINGO_LOG(INFO) << "vector backward scan query: " << tmp.ToString() << ", result:" << result.ToString();
     if (tmp.ok()) {
       CHECK_EQ(result.vectors[0].id, g_vector_ids[g_vector_ids.size() - 1]);
       CHECK_EQ(result.vectors[1].id, g_vector_ids[g_vector_ids.size() - 2]);
     }
   }
+
+  {
+    // forward with scalar filter
+    dingodb::sdk::ScanQueryParam param;
+    param.vector_id_start = g_vector_ids[0];
+    param.vector_id_end = g_vector_ids[g_vector_ids.size() - 1];
+    param.with_scalar_data = true;
+    param.use_scalar_filter = true;
+
+    int64_t filter_id = 5;
+    {
+      dingodb::sdk::ScalarValue scalar_value;
+      scalar_value.type = dingodb::sdk::ScalarFieldType::kInt64;
+
+      dingodb::sdk::ScalarField field;
+      field.long_data = filter_id;
+      scalar_value.fields.push_back(field);
+      param.scalar_data.insert(std::make_pair(g_scalar_col[0], scalar_value));
+    }
+
+    dingodb::sdk::ScanQueryResult result;
+    Status tmp;
+    if (use_index_name) {
+      tmp = g_vector_client->ScanQueryByIndexName(g_schema_id, g_index_name, param, result);
+    } else {
+      tmp = g_vector_client->ScanQueryByIndexId(g_index_id, param, result);
+    }
+
+    DINGO_LOG(INFO) << "vector forward scan query with filter:" << tmp.ToString() << ", result:" << result.ToString();
+    if (tmp.ok()) {
+      CHECK_EQ(result.vectors[0].id, filter_id);
+    }
+  }
 }
 
 static void VectorGetIndexMetrics(bool use_index_name = false) {

diff --git a/src/sdk/vector.h b/src/sdk/vector.h
@@ -391,9 +391,8 @@ struct ScanQueryParam {
   std::vector<std::string> selected_keys;
   bool with_table_data{false};  // Default false, if true, response without table data
 
-  // TODO: support use_scalar_filter
   bool use_scalar_filter{false};
-  // std::map<std::string, ScalarValue> scalar_data;
+  std::map<std::string, ScalarValue> scalar_data;
 
   explicit ScanQueryParam() = default;
 
@@ -406,7 +405,8 @@ struct ScanQueryParam {
         with_scalar_data(other.with_scalar_data),
         selected_keys(std::move(other.selected_keys)),
         with_table_data(other.with_table_data),
-        use_scalar_filter(other.use_scalar_filter) {}
+        use_scalar_filter(other.use_scalar_filter),
+        scalar_data(std::move(other.scalar_data)) {}
 
   ScanQueryParam& operator=(ScanQueryParam&& other) noexcept {
     if (this != &other) {
@@ -419,7 +419,7 @@ struct ScanQueryParam {
       selected_keys = std::move(other.selected_keys);
       with_table_data = other.with_table_data;
       use_scalar_filter = other.use_scalar_filter;
-      // You can add more fields here if you add more fields to the struct
+      scalar_data = std::move(other.scalar_data);
     }
     return *this;
   }

diff --git a/src/sdk/vector/vector_scan_query_task.cc b/src/sdk/vector/vector_scan_query_task.cc
@@ -183,7 +183,14 @@ void VectorScanQueryPartTask::FillVectorScanQueryRpcRequest(pb::index::VectorSca
 
   request->set_without_table_data(!scan_query_param_.with_table_data);
   request->set_use_scalar_filter(scan_query_param_.use_scalar_filter);
-  CHECK(!scan_query_param_.use_scalar_filter) << "not support scalar filter now";
+  if (scan_query_param_.use_scalar_filter) {
+    request->set_without_scalar_data(false);
+
+    for (const auto& [key, value] : scan_query_param_.scalar_data) {
+      request->mutable_scalar_for_filter()->mutable_scalar_data()->insert(
+          {key, ScalarValue2InternalScalarValuePB(value)});
+    }
+  }
 }
 
 void VectorScanQueryPartTask::VectorScanQueryRpcCallback(Status status, VectorScanQueryRpc* rpc) {