Push down the limit operator to segcore

Signed-off-by: longjiquan <[email protected]>
longjiquan · Jul 31, 2023 · 034e0c1 · 034e0c1
1 parent 7f51655
commit 034e0c1
Show file tree

Hide file tree

Showing 20 changed files with 581 additions and 273 deletions.
diff --git a/internal/core/CMakeLists.txt b/internal/core/CMakeLists.txt
@@ -123,8 +123,12 @@ if (LINUX OR MSYS)
                   "-DELPP_THREAD_SAFE"
                   "-fopenmp"
                   "-Werror"
-                  "-O3"
                   )
+    if (CMAKE_BUILD_TYPE STREQUAL "Release")
+        append_flags( CMAKE_CXX_FLAGS
+                "-O3"
+                )
+    endif()
 endif ()
 
 if ( APPLE )

diff --git a/internal/core/src/pb/plan.pb.cc b/internal/core/src/pb/plan.pb.cc
diff --git a/internal/core/src/pb/plan.pb.h b/internal/core/src/pb/plan.pb.h
diff --git a/internal/core/src/query/PlanNode.h b/internal/core/src/query/PlanNode.h
@@ -58,7 +58,8 @@ struct RetrievePlanNode : PlanNode {
     accept(PlanNodeVisitor&) override;
 
     std::optional<ExprPtr> predicate_;
-    bool is_count;
+    bool is_count_;
+    int64_t limit_;
 };
 
 }  // namespace milvus::query
diff --git a/internal/core/src/query/PlanProto.cpp b/internal/core/src/query/PlanProto.cpp
@@ -205,7 +205,7 @@ ProtoParser::RetrievePlanNodeFromProto(
     auto plan_node = [&]() -> std::unique_ptr<RetrievePlanNode> {
         auto node = std::make_unique<RetrievePlanNode>();
         if (plan_node_proto.has_predicates()) {  // version before 2023.03.30.
-            node->is_count = false;
+            node->is_count_ = false;
             auto& predicate_proto = plan_node_proto.predicates();
             auto expr_opt = [&]() -> ExprPtr {
                 return ParseExpr(predicate_proto);
@@ -220,7 +220,8 @@ ProtoParser::RetrievePlanNodeFromProto(
                 }();
                 node->predicate_ = std::move(expr_opt);
             }
-            node->is_count = query.is_count();
+            node->is_count_ = query.is_count();
+            node->limit_ = query.limit();
         }
         return node;
     }();

diff --git a/internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp b/internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp
@@ -148,20 +148,20 @@ ExecPlanNodeVisitor::visit(RetrievePlanNode& node) {
 
     auto active_count = segment->get_active_count(timestamp_);
 
-    if (active_count == 0 && !node.is_count) {
+    if (active_count == 0 && !node.is_count_) {
         retrieve_result_opt_ = std::move(retrieve_result);
         return;
     }
 
-    if (active_count == 0 && node.is_count) {
+    if (active_count == 0 && node.is_count_) {
         retrieve_result = *(wrap_num_entities(0));
         retrieve_result_opt_ = std::move(retrieve_result);
         return;
     }
 
     BitsetType bitset_holder;
     // For case that retrieve by expression, bitset will be allocated when expression is being executed.
-    if (node.is_count) {
+    if (node.is_count_) {
         bitset_holder.resize(active_count);
     }
 
@@ -176,27 +176,27 @@ ExecPlanNodeVisitor::visit(RetrievePlanNode& node) {
 
     segment->mask_with_delete(bitset_holder, active_count, timestamp_);
     // if bitset_holder is all 1's, we got empty result
-    if (bitset_holder.all() && !node.is_count) {
+    if (bitset_holder.all() && !node.is_count_) {
         retrieve_result_opt_ = std::move(retrieve_result);
         return;
     }
 
-    if (node.is_count) {
+    if (node.is_count_) {
         auto cnt = bitset_holder.size() - bitset_holder.count();
         retrieve_result = *(wrap_num_entities(cnt));
         retrieve_result_opt_ = std::move(retrieve_result);
         return;
     }
 
-    BitsetView final_view = bitset_holder;
-    auto seg_offsets =
-        GetExprUsePkIndex() && IsTermExpr(node.predicate_.value().get())
-            ? segment->search_ids(
-                  final_view, expr_cached_pk_id_offsets_, timestamp_)
-            : segment->search_ids(bitset_holder.flip(), timestamp_);
-    retrieve_result.result_offsets_.assign(
-        (int64_t*)seg_offsets.data(),
-        (int64_t*)seg_offsets.data() + seg_offsets.size());
+    bitset_holder.flip();
+    if (GetExprUsePkIndex() && IsTermExpr(node.predicate_.value().get())) {
+        segment->search_ids_filter(
+            bitset_holder, expr_cached_pk_id_offsets_, timestamp_);
+    } else {
+        segment->search_ids_filter(bitset_holder, timestamp_);
+    }
+    retrieve_result.result_offsets_ =
+        segment->find_first(node.limit_, bitset_holder);
     retrieve_result_opt_ = std::move(retrieve_result);
 }