From 03dffb4f58454e80a022733ad8a0c9fc781b1d0d Mon Sep 17 00:00:00 2001
From: Congqi Xia <congqi.xia@zilliz.com>
Date: Mon, 6 Jan 2025 11:57:02 +0800
Subject: [PATCH] enhance: Accelerate `find_first` by utilizing bitset simd
 methods

Related to #39003

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 .../src/segcore/ChunkedSegmentSealedImpl.cpp  | 18 ++++++++------
 internal/core/src/segcore/InsertRecord.h      | 24 +++++++++++--------
 .../core/src/segcore/SegmentSealedImpl.cpp    | 19 +++++++++------
 3 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
index b31fec905ea4b..0610f3adabd10 100644
--- a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
+++ b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
@@ -1255,20 +1255,24 @@ ChunkedSegmentSealedImpl::find_first(int64_t limit,
     std::vector<int64_t> seg_offsets;
     seg_offsets.reserve(limit);
 
+    // flip bitset since `find_next` is used to find true.
+    auto flipped = bitset.clone();
+    flipped.flip();
+
     int64_t offset = 0;
-    for (; hit_num < limit && offset < num_rows_.value(); offset++) {
+    std::optional<size_t> result = flipped.find_first();
+    while (result.has_value() && hit_num < limit) {
+        hit_num++;
+        seg_offsets.push_back(result.value());
+        offset = result.value();
         if (offset >= size) {
             // In fact, this case won't happen on sealed segments.
             continue;
         }
-
-        if (!bitset[offset]) {
-            seg_offsets.push_back(offset);
-            hit_num++;
-        }
+        result = flipped.find_next(offset);
     }
 
-    return {seg_offsets, more_hit_than_limit && offset != num_rows_.value()};
+    return {seg_offsets, more_hit_than_limit && result.has_value()};
 }
 
 ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl(
diff --git a/internal/core/src/segcore/InsertRecord.h b/internal/core/src/segcore/InsertRecord.h
index 3f8bb5a4d3738..59aa76cc315ee 100644
--- a/internal/core/src/segcore/InsertRecord.h
+++ b/internal/core/src/segcore/InsertRecord.h
@@ -252,20 +252,24 @@ class OffsetOrderedArray : public OffsetMap {
         limit = std::min(limit, cnt);
         std::vector<int64_t> seg_offsets;
         seg_offsets.reserve(limit);
-        auto it = array_.begin();
-        for (; hit_num < limit && it != array_.end(); it++) {
-            auto seg_offset = it->second;
-            if (seg_offset >= size) {
+        // flip bitset since `find_first` & `find_next` is used to find true.
+        // could be optimized by support find false in bitset.
+        auto flipped = bitset.clone();
+        flipped.flip();
+
+        int64_t offset = 0;
+        std::optional<size_t> result = flipped.find_first();
+        while (result.has_value() && hit_num < limit) {
+            hit_num++;
+            seg_offsets.push_back(result.value());
+            offset = result.value();
+            if (offset >= size) {
                 // In fact, this case won't happen on sealed segments.
                 continue;
             }
-
-            if (!bitset[seg_offset]) {
-                seg_offsets.push_back(seg_offset);
-                hit_num++;
-            }
+            result = flipped.find_next(offset);
         }
-        return {seg_offsets, more_hit_than_limit && it != array_.end()};
+        return {seg_offsets, more_hit_than_limit && result.has_value()};
     }
 
     void
diff --git a/internal/core/src/segcore/SegmentSealedImpl.cpp b/internal/core/src/segcore/SegmentSealedImpl.cpp
index ee8652da06d50..e751f04035724 100644
--- a/internal/core/src/segcore/SegmentSealedImpl.cpp
+++ b/internal/core/src/segcore/SegmentSealedImpl.cpp
@@ -1714,20 +1714,25 @@ SegmentSealedImpl::find_first(int64_t limit, const BitsetType& bitset) const {
     std::vector<int64_t> seg_offsets;
     seg_offsets.reserve(limit);
 
+    // flip bitset since `find_first` & `find_next` is used to find true.
+    // could be optimized by support find false in bitset.
+    auto flipped = bitset.clone();
+    flipped.flip();
+
     int64_t offset = 0;
-    for (; hit_num < limit && offset < num_rows_.value(); offset++) {
+    std::optional<size_t> result = flipped.find_first();
+    while (result.has_value() && hit_num < limit) {
+        hit_num++;
+        seg_offsets.push_back(result.value());
+        offset = result.value();
         if (offset >= size) {
             // In fact, this case won't happen on sealed segments.
             continue;
         }
-
-        if (!bitset[offset]) {
-            seg_offsets.push_back(offset);
-            hit_num++;
-        }
+        result = flipped.find_next(offset);
     }
 
-    return {seg_offsets, more_hit_than_limit && offset != num_rows_.value()};
+    return {seg_offsets, more_hit_than_limit && result.has_value()};
 }
 
 SegcoreError