From 72f5b85c0591474fd5f427194b869aa3470595c1 Mon Sep 17 00:00:00 2001 From: congqixia Date: Tue, 7 Jan 2025 10:34:54 +0800 Subject: [PATCH] enhance: Accelerate `find_first` by utilizing bitset simd methods (#39004) Related to #39003 Signed-off-by: Congqi Xia --- .../src/segcore/ChunkedSegmentSealedImpl.cpp | 18 +++++++++++------- .../core/src/segcore/SegmentSealedImpl.cpp | 19 ++++++++++++------- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp index b31fec905e..0610f3adab 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 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 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/SegmentSealedImpl.cpp b/internal/core/src/segcore/SegmentSealedImpl.cpp index ee8652da06..e751f04035 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 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 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