Enable search for sealed

Signed-off-by: FluorineDog <guilin.gou@zilliz.com>

docs/developer_guides/chap01_system_overview.md

@@ -78,9 +78,9 @@ In order to boost throughput, we model Milvus as a stream-driven system.
 
 ```go
 type Component interface {
-  Init() error
-  Start() error
-  Stop() error
+  Init()
+  Start()
+  Stop()
   GetComponentStates() (ComponentStates, error)
   GetTimeTickChannel() (string, error)
   GetStatisticsChannel() (string, error)

internal/core/src/query/Plan.cpp

@@ -209,7 +209,7 @@ Parser::ParseVecNode(const Json& out_body) {
         }
     }();
     vec_node->query_info_.topK_ = topK;
-    vec_node->query_info_.metric_type_ = vec_info.at("metric_type");
+    vec_node->query_info_.metric_type_ = GetMetricType(vec_info.at("metric_type"));
     vec_node->query_info_.search_params_ = vec_info.at("params");
     vec_node->query_info_.field_offset_ = field_offset;
     vec_node->placeholder_tag_ = vec_info.at("query");

internal/core/src/query/PlanNode.h

@@ -41,7 +41,8 @@ using PlanNodePtr = std::unique_ptr<PlanNode>;
 struct QueryInfo {
     int64_t topK_;
     FieldOffset field_offset_;
-    std::string metric_type_;  // TODO: use enum
+    MetricType metric_type_;
+    std::string deprecated_metric_type_;  // TODO: use enum
     nlohmann::json search_params_;
 };
 

internal/core/src/query/SearchOnGrowing.cpp

@@ -63,7 +63,7 @@ FloatSearch(const segcore::SegmentGrowingImpl& segment,
     auto dim = field.get_dim();
     auto topK = info.topK_;
     auto total_count = topK * num_queries;
-    auto metric_type = GetMetricType(info.metric_type_);
+    auto metric_type = info.metric_type_;
 
     // step 3: small indexing search
     // std::vector<int64_t> final_uids(total_count, -1);
@@ -138,7 +138,7 @@ BinarySearch(const segcore::SegmentGrowingImpl& segment,
     auto& record = segment.get_insert_record();
     // step 1: binary search to find the barrier of the snapshot
     // auto ins_barrier = get_barrier(record, timestamp);
-    auto metric_type = GetMetricType(info.metric_type_);
+    auto metric_type = info.metric_type_;
     // auto del_barrier = get_barrier(deleted_record_, timestamp);
 
 #if 0
@@ -195,37 +195,24 @@ BinarySearch(const segcore::SegmentGrowingImpl& segment,
 }
 
 // TODO: refactor and merge this into one
-template <typename VectorType>
 void
 SearchOnGrowing(const segcore::SegmentGrowingImpl& segment,
                 int64_t ins_barrier,
                 const query::QueryInfo& info,
-                const EmbeddedType<VectorType>* query_data,
+                const void* query_data,
                 int64_t num_queries,
                 const faiss::BitsetView& bitset,
                 QueryResult& results) {
-    static_assert(IsVector<VectorType>);
-    if constexpr (std::is_same_v<VectorType, FloatVector>) {
-        FloatSearch(segment, info, query_data, num_queries, ins_barrier, bitset, results);
+    // TODO: add data_type to info
+    auto data_type = segment.get_schema()[info.field_offset_].get_data_type();
+    Assert(datatype_is_vector(data_type));
+    if (data_type == DataType::VECTOR_FLOAT) {
+        auto typed_data = reinterpret_cast<const float*>(query_data);
+        FloatSearch(segment, info, typed_data, num_queries, ins_barrier, bitset, results);
     } else {
-        BinarySearch(segment, info, query_data, num_queries, ins_barrier, bitset, results);
+        auto typed_data = reinterpret_cast<const uint8_t*>(query_data);
+        BinarySearch(segment, info, typed_data, num_queries, ins_barrier, bitset, results);
     }
 }
-template void
-SearchOnGrowing<FloatVector>(const segcore::SegmentGrowingImpl& segment,
-                             int64_t ins_barrier,
-                             const query::QueryInfo& info,
-                             const EmbeddedType<FloatVector>* query_data,
-                             int64_t num_queries,
-                             const faiss::BitsetView& bitset,
-                             QueryResult& results);
-template void
-SearchOnGrowing<BinaryVector>(const segcore::SegmentGrowingImpl& segment,
-                              int64_t ins_barrier,
-                              const query::QueryInfo& info,
-                              const EmbeddedType<BinaryVector>* query_data,
-                              int64_t num_queries,
-                              const faiss::BitsetView& bitset,
-                              QueryResult& results);
 
 }  // namespace milvus::query

internal/core/src/query/SearchOnGrowing.h

@@ -20,12 +20,11 @@ namespace milvus::query {
 using BitmapChunk = boost::dynamic_bitset<>;
 using BitmapSimple = std::deque<BitmapChunk>;
 
-template <typename VectorType>
 void
 SearchOnGrowing(const segcore::SegmentGrowingImpl& segment,
                 int64_t ins_barrier,
                 const query::QueryInfo& info,
-                const EmbeddedType<VectorType>* query_data,
+                const void* query_data,
                 int64_t num_queries,
                 const faiss::BitsetView& bitset,
                 QueryResult& results);

internal/core/src/query/SearchOnSealed.cpp

@@ -63,9 +63,7 @@ SearchOnSealed(const Schema& schema,
 
     Assert(record.is_ready(field_offset));
     auto indexing_entry = record.get_entry(field_offset);
-    std::cout << " SearchOnSealed, indexing_entry->metric:" << indexing_entry->metric_type_ << std::endl;
-    std::cout << " SearchOnSealed, query_info.metric_type_:" << query_info.metric_type_ << std::endl;
-    Assert(indexing_entry->metric_type_ == GetMetricType(query_info.metric_type_));
+    Assert(indexing_entry->metric_type_ == query_info.metric_type_);
 
     auto final = [&] {
         auto ds = knowhere::GenDataset(num_queries, dim, query_data);

internal/core/src/query/generated/ExecExprVisitor.h

@@ -37,7 +37,7 @@ class ExecExprVisitor : public ExprVisitor {
 
  public:
     using RetType = std::deque<boost::dynamic_bitset<>>;
-    ExecExprVisitor(const segcore::SegmentGrowingImpl& segment, int64_t row_count)
+    ExecExprVisitor(const segcore::SegmentInternalInterface& segment, int64_t row_count)
         : segment_(segment), row_count_(row_count) {
     }
     RetType

internal/core/src/query/generated/ExecPlanNodeVisitor.h

@@ -29,7 +29,7 @@ class ExecPlanNodeVisitor : public PlanNodeVisitor {
 
  public:
     using RetType = QueryResult;
-    ExecPlanNodeVisitor(const segcore::SegmentGrowing& segment,
+    ExecPlanNodeVisitor(const segcore::SegmentInterface& segment,
                         Timestamp timestamp,
                         const PlaceholderGroup& placeholder_group)
         : segment_(segment), timestamp_(timestamp), placeholder_group_(placeholder_group) {
@@ -53,7 +53,7 @@ class ExecPlanNodeVisitor : public PlanNodeVisitor {
 
  private:
     // std::optional<RetType> ret_;
-    const segcore::SegmentGrowing& segment_;
+    const segcore::SegmentInterface& segment_;
     Timestamp timestamp_;
     const PlaceholderGroup& placeholder_group_;
 

internal/core/src/query/visitors/ExecExprVisitor.cpp

@@ -25,7 +25,7 @@ namespace impl {
 class ExecExprVisitor : ExprVisitor {
  public:
     using RetType = std::deque<boost::dynamic_bitset<>>;
-    ExecExprVisitor(const segcore::SegmentGrowingImpl& segment, int64_t row_count)
+    ExecExprVisitor(const segcore::SegmentInternalInterface& segment, int64_t row_count)
         : segment_(segment), row_count_(row_count) {
     }
     RetType

internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp

@@ -28,7 +28,7 @@ namespace impl {
 class ExecPlanNodeVisitor : PlanNodeVisitor {
  public:
     using RetType = QueryResult;
-    ExecPlanNodeVisitor(const segcore::SegmentGrowing& segment,
+    ExecPlanNodeVisitor(const segcore::SegmentInterface& segment,
                         Timestamp timestamp,
                         const PlaceholderGroup& placeholder_group)
         : segment_(segment), timestamp_(timestamp), placeholder_group_(placeholder_group) {
@@ -52,7 +52,7 @@ class ExecPlanNodeVisitor : PlanNodeVisitor {
 
  private:
     // std::optional<RetType> ret_;
-    const segcore::SegmentGrowing& segment_;
+    const segcore::SegmentInterface& segment_;
     Timestamp timestamp_;
     const PlaceholderGroup& placeholder_group_;
 
@@ -66,7 +66,7 @@ void
 ExecPlanNodeVisitor::VectorVisitorImpl(VectorPlanNode& node) {
     // TODO: optimize here, remove the dynamic cast
     assert(!ret_.has_value());
-    auto segment = dynamic_cast<const segcore::SegmentGrowingImpl*>(&segment_);
+    auto segment = dynamic_cast<const segcore::SegmentInternalInterface*>(&segment_);
     AssertInfo(segment, "support SegmentSmallIndex Only");
     RetType ret;
     auto& ph = placeholder_group_.at(0);
@@ -76,7 +76,7 @@ ExecPlanNodeVisitor::VectorVisitorImpl(VectorPlanNode& node) {
     aligned_vector<uint8_t> bitset_holder;
     BitsetView view;
     // TODO: add API to unify row_count
-    auto row_count = segcore::get_barrier(segment->get_insert_record(), timestamp_);
+    auto row_count = segment->get_row_count();
 
     if (node.predicate_.has_value()) {
         ExecExprVisitor::RetType expr_ret = ExecExprVisitor(*segment, row_count).call_child(*node.predicate_.value());
@@ -84,12 +84,7 @@ ExecPlanNodeVisitor::VectorVisitorImpl(VectorPlanNode& node) {
         view = BitsetView(bitset_holder.data(), bitset_holder.size() * 8);
     }
 
-    auto& sealed_indexing = segment->get_sealed_indexing_record();
-    if (sealed_indexing.is_ready(node.query_info_.field_offset_)) {
-        SearchOnSealed(segment->get_schema(), sealed_indexing, node.query_info_, src_data, num_queries, view, ret);
-    } else {
-        SearchOnGrowing<VectorType>(*segment, row_count, node.query_info_, src_data, num_queries, view, ret);
-    }
+    segment->vector_search(row_count, node.query_info_, src_data, num_queries, view, ret);
 
     ret_ = ret;
 }

internal/core/src/query/visitors/ShowPlanNodeVisitor.cpp

@@ -54,12 +54,12 @@ ShowPlanNodeVisitor::visit(FloatVectorANNS& node) {
     assert(!ret_);
     auto& info = node.query_info_;
     Json json_body{
-        {"node_type", "FloatVectorANNS"},             //
-        {"metric_type", info.metric_type_},           //
-        {"field_offset_", info.field_offset_.get()},  //
-        {"topK", info.topK_},                         //
-        {"search_params", info.search_params_},       //
-        {"placeholder_tag", node.placeholder_tag_},   //
+        {"node_type", "FloatVectorANNS"},                      //
+        {"metric_type", MetricTypeToName(info.metric_type_)},  //
+        {"field_offset_", info.field_offset_.get()},           //
+        {"topK", info.topK_},                                  //
+        {"search_params", info.search_params_},                //
+        {"placeholder_tag", node.placeholder_tag_},            //
     };
     if (node.predicate_.has_value()) {
         ShowExprVisitor expr_show;
@@ -76,12 +76,12 @@ ShowPlanNodeVisitor::visit(BinaryVectorANNS& node) {
     assert(!ret_);
     auto& info = node.query_info_;
     Json json_body{
-        {"node_type", "BinaryVectorANNS"},            //
-        {"metric_type", info.metric_type_},           //
-        {"field_offset_", info.field_offset_.get()},  //
-        {"topK", info.topK_},                         //
-        {"search_params", info.search_params_},       //
-        {"placeholder_tag", node.placeholder_tag_},   //
+        {"node_type", "BinaryVectorANNS"},                     //
+        {"metric_type", MetricTypeToName(info.metric_type_)},  //
+        {"field_offset_", info.field_offset_.get()},           //
+        {"topK", info.topK_},                                  //
+        {"search_params", info.search_params_},                //
+        {"placeholder_tag", node.placeholder_tag_},            //
     };
     if (node.predicate_.has_value()) {
         ShowExprVisitor expr_show;

internal/core/src/segcore/SegmentGrowingImpl.cpp

@@ -19,6 +19,7 @@
 #include <knowhere/index/vector_index/adapter/VectorAdapter.h>
 #include <knowhere/index/vector_index/VecIndexFactory.h>
 #include <faiss/utils/distances.h>
+#include <query/SearchOnSealed.h>
 #include "query/generated/ExecPlanNodeVisitor.h"
 #include "segcore/SegmentGrowingImpl.h"
 #include "query/PlanNode.h"
@@ -237,17 +238,6 @@ SegmentGrowingImpl::GetMemoryUsageInBytes() const {
     return total_bytes;
 }
 
-QueryResult
-SegmentGrowingImpl::Search(const query::Plan* plan,
-                           const query::PlaceholderGroup** placeholder_groups,
-                           const Timestamp* timestamps,
-                           int64_t num_groups) const {
-    Assert(num_groups == 1);
-    query::ExecPlanNodeVisitor visitor(*this, timestamps[0], *placeholder_groups[0]);
-    auto results = visitor.get_moved_result(*plan->plan_node_);
-    return results;
-}
-
 Status
 SegmentGrowingImpl::LoadIndexing(const LoadIndexInfo& info) {
     auto field_offset = schema_->get_offset(FieldName(info.field_name));
@@ -270,5 +260,19 @@ SegmentGrowingImpl::num_chunk_data() const {
     auto size = get_insert_record().ack_responder_.GetAck();
     return upper_div(size, chunk_size_);
 }
+void
+SegmentGrowingImpl::vector_search(int64_t vec_count,
+                                  query::QueryInfo query_info,
+                                  const void* query_data,
+                                  int64_t query_count,
+                                  const BitsetView& bitset,
+                                  QueryResult& output) const {
+    auto& sealed_indexing = this->get_sealed_indexing_record();
+    if (sealed_indexing.is_ready(query_info.field_offset_)) {
+        query::SearchOnSealed(this->get_schema(), sealed_indexing, query_info, query_data, query_count, bitset, output);
+    } else {
+        SearchOnGrowing(*this, vec_count, query_info, query_data, query_count, bitset, output);
+    }
+}
 
 }  // namespace milvus::segcore

internal/core/src/segcore/SegmentGrowingImpl.h

@@ -53,12 +53,6 @@ class SegmentGrowingImpl : public SegmentGrowing {
     Status
     Delete(int64_t reserverd_offset, int64_t size, const int64_t* row_ids, const Timestamp* timestamps) override;
 
-    QueryResult
-    Search(const query::Plan* Plan,
-           const query::PlaceholderGroup* placeholder_groups[],
-           const Timestamp timestamps[],
-           int64_t num_groups) const override;
-
     // stop receive insert requests
     // will move data to immutable vector or something
     Status
@@ -181,6 +175,14 @@ class SegmentGrowingImpl : public SegmentGrowing {
           indexing_record_(*schema_, chunk_size) {
     }
 
+    void
+    vector_search(int64_t vec_count,
+                  query::QueryInfo query_info,
+                  const void* query_data,
+                  int64_t query_count,
+                  const BitsetView& bitset,
+                  QueryResult& output) const override;
+
  public:
     std::shared_ptr<DeletedRecord::TmpBitmap>
     get_deleted_bitmap(int64_t del_barrier, Timestamp query_timestamp, int64_t insert_barrier, bool force = false);

internal/core/src/segcore/SegmentInterface.cpp

@@ -10,6 +10,7 @@
 // or implied. See the License for the specific language governing permissions and limitations under the License
 
 #include "segcore/SegmentInterface.h"
+#include "query/generated/ExecPlanNodeVisitor.h"
 namespace milvus::segcore {
 class Naive;
 
@@ -39,4 +40,16 @@ SegmentInterface::FillTargetEntry(const query::Plan* plan, QueryResult& results)
         results.row_data_.emplace_back(std::move(blob));
     }
 }
+
+QueryResult
+SegmentInterface::Search(const query::Plan* plan,
+                         const query::PlaceholderGroup** placeholder_groups,
+                         const Timestamp* timestamps,
+                         int64_t num_groups) const {
+    Assert(num_groups == 1);
+    query::ExecPlanNodeVisitor visitor(*this, timestamps[0], *placeholder_groups[0]);
+    auto results = visitor.get_moved_result(*plan->plan_node_);
+    return results;
+}
+
 }  // namespace milvus::segcore

internal/core/src/segcore/SegmentInterface.h

@@ -17,6 +17,7 @@
 #include "IndexingEntry.h"
 #include <knowhere/index/vector_index/VecIndex.h>
 #include "common/SystemProperty.h"
+#include "query/PlanNode.h"
 
 namespace milvus::segcore {
 
@@ -25,11 +26,11 @@ class SegmentInterface {
     void
     FillTargetEntry(const query::Plan* plan, QueryResult& results) const;
 
-    virtual QueryResult
+    QueryResult
     Search(const query::Plan* Plan,
            const query::PlaceholderGroup* placeholder_groups[],
            const Timestamp timestamps[],
-           int64_t num_groups) const = 0;
+           int64_t num_groups) const;
 
     virtual int64_t
     GetMemoryUsageInBytes() const = 0;
@@ -71,6 +72,14 @@ class SegmentInternalInterface : public SegmentInterface {
     }
 
  public:
+    virtual void
+    vector_search(int64_t vec_count,
+                  query::QueryInfo query_info,
+                  const void* query_data,
+                  int64_t query_count,
+                  const BitsetView& bitset,
+                  QueryResult& output) const = 0;
+
     virtual int64_t
     num_chunk_index_safe(FieldOffset field_offset) const = 0;
 

internal/core/src/segcore/SegmentSealedImpl.cpp

@@ -10,6 +10,7 @@
 // or implied. See the License for the specific language governing permissions and limitations under the License
 
 #include "segcore/SegmentSealedImpl.h"
+#include "query/SearchOnSealed.h"
 namespace milvus::segcore {
 void
 SegmentSealedImpl::LoadIndex(const LoadIndexInfo& info) {
@@ -107,14 +108,6 @@ SegmentSealedImpl::chunk_index_impl(FieldOffset field_offset, int64_t chunk_id)
     return nullptr;
 }
 
-QueryResult
-SegmentSealedImpl::Search(const query::Plan* Plan,
-                          const query::PlaceholderGroup** placeholder_groups,
-                          const Timestamp* timestamps,
-                          int64_t num_groups) const {
-    PanicInfo("unimplemented");
-}
-
 int64_t
 SegmentSealedImpl::GetMemoryUsageInBytes() const {
     // TODO: add estimate for index
@@ -134,6 +127,20 @@ SegmentSealedImpl::get_schema() const {
     return *schema_;
 }
 
+void
+SegmentSealedImpl::vector_search(int64_t vec_count,
+                                 query::QueryInfo query_info,
+                                 const void* query_data,
+                                 int64_t query_count,
+                                 const BitsetView& bitset,
+                                 QueryResult& output) const {
+    auto field_offset = query_info.field_offset_;
+    auto& field_meta = schema_->operator[](field_offset);
+    Assert(field_meta.is_vector());
+    Assert(vec_indexings_.is_ready(field_offset));
+    query::SearchOnSealed(*schema_, vec_indexings_, query_info, query_data, query_count, bitset, output);
+}
+
 SegmentSealedPtr
 CreateSealedSegment(SchemaPtr schema, int64_t chunk_size) {
     return std::make_unique<SegmentSealedImpl>(schema);

internal/core/src/segcore/SegmentSealedImpl.h

@@ -26,12 +26,6 @@ class SegmentSealedImpl : public SegmentSealed {
     LoadFieldData(const LoadFieldDataInfo& info) override;
 
  public:
-    QueryResult
-    Search(const query::Plan* Plan,
-           const query::PlaceholderGroup* placeholder_groups[],
-           const Timestamp timestamps[],
-           int64_t num_groups) const override;
-
     int64_t
     GetMemoryUsageInBytes() const override;
 
@@ -104,6 +98,14 @@ class SegmentSealedImpl : public SegmentSealed {
         }
     }
 
+    void
+    vector_search(int64_t vec_count,
+                  query::QueryInfo query_info,
+                  const void* query_data,
+                  int64_t query_count,
+                  const BitsetView& bitset,
+                  QueryResult& output) const override;
+
     bool
     is_all_ready() const {
         // TODO: optimize here

internal/core/src/segcore/plan_c.cpp

@@ -78,7 +78,8 @@ GetTopK(CPlan plan) {
 const char*
 GetMetricType(CPlan plan) {
     auto query_plan = static_cast<milvus::query::Plan*>(plan);
-    return strdup(query_plan->plan_node_->query_info_.metric_type_.c_str());
+    auto metric_str = milvus::MetricTypeToName(query_plan->plan_node_->query_info_.metric_type_);
+    return strdup(metric_str.c_str());
 }
 
 void

internal/core/src/segcore/segment_c.cpp

@@ -20,6 +20,7 @@
 #include "common/type_c.h"
 #include <knowhere/index/vector_index/VecIndex.h>
 #include <knowhere/index/vector_index/adapter/VectorAdapter.h>
+#include "common/Types.h"
 
 //////////////////////////////    common interfaces    //////////////////////////////
 CSegmentInterface
@@ -79,7 +80,7 @@ Search(CSegmentInterface c_segment,
     auto status = CStatus();
     try {
         *query_result = segment->Search(plan, placeholder_groups.data(), timestamps, num_groups);
-        if (plan->plan_node_->query_info_.metric_type_ != "IP") {
+        if (plan->plan_node_->query_info_.metric_type_ != milvus::MetricType::METRIC_INNER_PRODUCT) {
             for (auto& dis : query_result->result_distances_) {
                 dis *= -1;
             }

internal/core/unittest/test_expr.cpp

@@ -235,7 +235,8 @@ TEST(Expr, ShowExecutor) {
     int64_t num_queries = 100L;
     auto raw_data = DataGen(schema, num_queries);
     auto& info = node->query_info_;
-    info.metric_type_ = "L2";
+
+    info.metric_type_ = MetricType::METRIC_L2;
     info.topK_ = 20;
     info.field_offset_ = FieldOffset(0);
     node->predicate_ = std::nullopt;