Support Roaring-backed bitsets

CMakeLists.txt

@@ -103,6 +103,8 @@ endif()
 find_package(xxHash REQUIRED)
 include_directories(${xxHash_INCLUDE_DIRS})
 find_package(simde REQUIRED)
+find_package(roaring REQUIRED)
+include_directories(${roaring_INCLUDE_DIRS})
 
 if(NOT WITH_LIGHT)
   find_package(opentelemetry-cpp REQUIRED)
@@ -114,6 +116,7 @@ if(NOT WITH_LIGHT)
 endif()
 
 if(APPLE)
+  add_definitions(-D_DARWIN_C_SOURCE)
   # Prevent SDK sysroot include dir from appearing as -I/-isystem in compile
   # commands. Transitive find_package calls (CURL, ZLIB, etc.) resolve headers
   # to the SDK sysroot, and the explicit -I flag breaks C++ wrapper headers
@@ -125,7 +128,7 @@ if(APPLE)
 endif()
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_OSX_DEPLOYMENT_TARGET
-    "10.15"
+    "11.0"
     CACHE STRING "Minimum OS X deployment version" FORCE)
 
 if(OPENMP_FOUND)
@@ -139,6 +142,7 @@ include(cmake/utils/compile_flags.cmake)
 find_package(milvus-common REQUIRED)
 include(cmake/libs/libhnsw.cmake)
 include(cmake/libs/libfaiss.cmake)
+target_link_libraries(faiss PUBLIC roaring::roaring)
 
 if(WITH_COVERAGE)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage")
@@ -268,6 +272,7 @@ list(APPEND KNOWHERE_LINKER_LIBS fmt::fmt)
 list(APPEND KNOWHERE_LINKER_LIBS Folly::folly)
 list(APPEND KNOWHERE_LINKER_LIBS milvus-common::milvus-common)
 list(APPEND KNOWHERE_LINKER_LIBS simde::simde)
+list(APPEND KNOWHERE_LINKER_LIBS roaring::roaring)
 
 add_library(knowhere SHARED ${KNOWHERE_SRCS})
 add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS})

cmake/libs/libdiskann.cmake

@@ -29,7 +29,7 @@ set(DISKANN_SOURCES
 
 find_package(folly REQUIRED)
 set(DISKANN_LINKER_LIBS PUBLIC ${AIO_LIBRARIES} ${DISKANN_BOOST_PROGRAM_OPTIONS_LIB} nlohmann_json::nlohmann_json
-         	Folly::folly fmt::fmt-header-only prometheus-cpp::core prometheus-cpp::push glog::glog)
+            Folly::folly fmt::fmt-header-only prometheus-cpp::core prometheus-cpp::push glog::glog roaring::roaring)
 if (WITH_CUVS)
     list(APPEND DISKANN_LINKER_LIBS PRIVATE cuvs::cuvs)
     list(APPEND DISKANN_SOURCES thirdparty/DiskANN/src/diskann_gpu.cpp)
@@ -49,7 +49,8 @@ target_link_libraries(
          fmt::fmt
          prometheus-cpp::core
          prometheus-cpp::push
-         glog::glog)
+         glog::glog
+         roaring::roaring)
 if(__X86_64)
   target_compile_options(
     diskann PRIVATE -fno-builtin-malloc -fno-builtin-calloc

cmake/libs/libfaiss.cmake

@@ -384,7 +384,7 @@ if(__X86_64)
                                             -mavx2 -mfma -mf16c -mpopcnt>)
   target_compile_definitions(faiss_avx2 PRIVATE COMPILE_SIMD_AVX2)
   target_include_directories(faiss_avx2 PRIVATE ${Boost_INCLUDE_DIRS})
-  target_link_libraries(faiss_avx2 PRIVATE milvus-common::milvus-common)
+  target_link_libraries(faiss_avx2 PRIVATE milvus-common::milvus-common roaring::roaring)
   add_library(faiss_avx512 OBJECT ${FAISS_AVX512_SRCS})
   target_compile_options(
     faiss_avx512
@@ -400,7 +400,7 @@ if(__X86_64)
             -mpopcnt>)
   target_compile_definitions(faiss_avx512 PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512)
   target_include_directories(faiss_avx512 PRIVATE ${Boost_INCLUDE_DIRS})
-  target_link_libraries(faiss_avx512 PRIVATE milvus-common::milvus-common)
+  target_link_libraries(faiss_avx512 PRIVATE milvus-common::milvus-common roaring::roaring)
 
   if(FAISS_ENABLE_AVX512_SPR)
     add_library(faiss_avx512_spr OBJECT ${FAISS_AVX512_SPR_SRCS})
@@ -424,7 +424,7 @@ if(__X86_64)
     target_compile_definitions(faiss_avx512_spr PRIVATE
                                COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR)
     target_include_directories(faiss_avx512_spr PRIVATE ${Boost_INCLUDE_DIRS})
-    target_link_libraries(faiss_avx512_spr PRIVATE milvus-common::milvus-common)
+    target_link_libraries(faiss_avx512_spr PRIVATE milvus-common::milvus-common roaring::roaring)
   endif()
 
   add_library(faiss STATIC ${FAISS_SRCS})

conanfile.py

@@ -137,6 +137,7 @@ def requirements(self):
         self.requires("libcurl/8.10.1#a3113369c86086b0e84231844e7ed0a9", force=True, override=True)
         self.requires("simde/0.8.2#5e1edfd5cba92f25d79bf6ef4616b972")
         self.requires("xxhash/0.8.3#caa6d0af1b951c247922e38fbcebdbe6")
+        self.requires("roaring/3.0.0")
         if self.settings.os == "Linux":
             self.requires("openblas/0.3.30")
         if not self.options.with_light:
@@ -283,6 +284,7 @@ def package_info(self):
             "milvus-common::milvus-common",
             "prometheus-cpp::core",
             "prometheus-cpp::push",
+            "roaring::roaring",
         ]
 
         self.cpp_info.components["libknowhere"].set_property(

include/knowhere/bitsetview.h

@@ -12,11 +12,15 @@
 #ifndef BITSET_H
 #define BITSET_H
 
+#include <roaring/roaring.h>
+
 #include <cassert>
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <sstream>
 #include <string>
+#include <vector>
 
 namespace knowhere {
 class BitsetView {
@@ -25,12 +29,38 @@ class BitsetView {
     ~BitsetView() = default;
 
     BitsetView(const uint8_t* data, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0)
-        : bits_(data), num_bits_(num_bits), num_filtered_out_bits_(num_filtered_out_bits), id_offset_(id_offset) {
+        : kind_(Kind::Dense),
+          bits_(data),
+          num_bits_(num_bits),
+          num_filtered_out_bits_(num_filtered_out_bits),
+          id_offset_(id_offset) {
+    }
+
+    BitsetView(const roaring_bitmap_t* bitmap, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0)
+        : kind_(Kind::Roaring),
+          roaring_(bitmap),
+          num_bits_(num_bits),
+          num_filtered_out_bits_(num_filtered_out_bits),
+          id_offset_(id_offset) {
     }
 
     BitsetView(const std::nullptr_t) : BitsetView() {
     }
 
+    BitsetView(const std::nullptr_t, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0)
+        : BitsetView(static_cast<const uint8_t*>(nullptr), num_bits, num_filtered_out_bits, id_offset) {
+    }
+
+    static BitsetView
+    FromFrozenRoaring(const void* data, size_t byte_size, size_t num_bits, size_t num_filtered_out_bits = 0,
+                      size_t id_offset = 0) {
+        const auto* bitmap = roaring_bitmap_frozen_view(static_cast<const char*>(data), byte_size);
+        BitsetView bitset(bitmap, num_bits, num_filtered_out_bits, id_offset);
+        bitset.owned_roaring_ = std::shared_ptr<const roaring_bitmap_t>(
+            bitmap, [](const roaring_bitmap_t* p) { roaring_bitmap_free(const_cast<roaring_bitmap_t*>(p)); });
+        return bitset;
+    }
+
     bool
     empty() const {
         return num_bits_ == 0;
@@ -54,6 +84,15 @@ class BitsetView {
         return num_filtered_out_bits_;
     }
 
+    void
+    set_count(size_t num_filtered_out_bits) {
+        if (out_ids_ != nullptr) {
+            num_filtered_out_ids_ = num_filtered_out_bits;
+            return;
+        }
+        num_filtered_out_bits_ = num_filtered_out_bits;
+    }
+
     size_t
     byte_size() const {
         return (num_bits_ + 8 - 1) >> 3;
@@ -64,6 +103,42 @@ class BitsetView {
         return bits_;
     }
 
+    bool
+    is_dense() const {
+        return kind_ == Kind::Dense;
+    }
+
+    bool
+    is_roaring() const {
+        return kind_ == Kind::Roaring;
+    }
+
+    const roaring_bitmap_t*
+    roaring() const {
+        return roaring_;
+    }
+
+    size_t
+    id_offset() const {
+        return id_offset_;
+    }
+
+    bool
+    can_iterate_roaring_without_mapping() const {
+        return kind_ == Kind::Roaring && out_ids_ == nullptr;
+    }
+
+    std::vector<uint8_t>
+    ToDense() const {
+        std::vector<uint8_t> dense(byte_size(), 0);
+        for (size_t i = 0; i < num_bits_; ++i) {
+            if (test(i)) {
+                dense[i >> 3] |= 0x1 << (i & 0x7);
+            }
+        }
+        return dense;
+    }
+
     bool
     has_out_ids() const {
         return out_ids_ != nullptr;
@@ -103,7 +178,13 @@ class BitsetView {
             out_id = out_ids_[out_id];
         }
         // when index is larger than the max_offset, ignore it
-        return (out_id >= static_cast<int64_t>(num_bits_)) || (bits_[out_id >> 3] & (0x1 << (out_id & 0x7)));
+        if (out_id >= static_cast<int64_t>(num_bits_)) {
+            return true;
+        }
+        if (kind_ == Kind::Roaring) {
+            return roaring_bitmap_contains(roaring_, static_cast<uint32_t>(out_id));
+        }
+        return bits_[out_id >> 3] & (0x1 << (out_id & 0x7));
     }
     // return the filtered ratio. if with id mapping, calculated by internal_ids rather than bits.
     float
@@ -126,6 +207,9 @@ class BitsetView {
             }
             return count;
         }
+        if (kind_ == Kind::Roaring) {
+            return roaring_bitmap_get_cardinality(roaring_);
+        }
         // if without id mapping, use a better algorithm to calculate the number of filtered out bits.
         size_t ret = 0;
         auto len_uint8 = byte_size();
@@ -166,6 +250,14 @@ class BitsetView {
             }
             return num_internal_ids_;
         }
+        if (kind_ == Kind::Roaring) {
+            for (size_t i = 0; i < num_bits_; i++) {
+                if (!test(i)) {
+                    return i;
+                }
+            }
+            return num_bits_;
+        }
         // if without id mapping, use a better algorithm to find the first valid index.
         size_t ret = 0;
         auto len_uint8 = byte_size();
@@ -211,7 +303,12 @@ class BitsetView {
     }
 
  private:
+    enum class Kind { Dense, Roaring };
+
+    Kind kind_ = Kind::Dense;
     const uint8_t* bits_ = nullptr;
+    const roaring_bitmap_t* roaring_ = nullptr;
+    std::shared_ptr<const roaring_bitmap_t> owned_roaring_;
     size_t num_bits_ = 0;
     size_t num_filtered_out_bits_ = 0;
 

src/index/index.cc

@@ -154,7 +154,8 @@ Index<T>::Search(const DataSetPtr dataset, const Json& json, const BitsetView& b
         if (bitset_.count() == 0) {
             // traverse bitset to get the filtered out num
             auto filtered_out_num = bitset_.get_filtered_out_num_();
-            bitset = BitsetView(bitset_.data(), bitset_.size(), filtered_out_num);
+            bitset = bitset_;
+            bitset.set_count(filtered_out_num);
         } else {
             // if bitset has filtered out num, use it
             bitset = bitset_;
@@ -212,7 +213,10 @@ Index<T>::AnnIterator(const DataSetPtr dataset, const Json& json, const BitsetVi
             return expected<std::vector<std::shared_ptr<IndexNode::iterator>>>::Err(Status::invalid_args, msg);
         }
 
-        const auto bitset = BitsetView(bitset_.data(), bitset_.size(), bitset_.get_filtered_out_num_());
+        auto bitset = bitset_;
+        if (bitset.count() == 0) {
+            bitset.set_count(bitset.get_filtered_out_num_());
+        }
 
 #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT)
         // note that this time includes only the initial search phase of iterator.
@@ -252,7 +256,10 @@ Index<T>::RangeSearch(const DataSetPtr dataset, const Json& json, const BitsetVi
             return expected<DataSetPtr>::Err(Status::invalid_args, msg);
         }
 
-        const auto bitset = BitsetView(bitset_.data(), bitset_.size(), bitset_.get_filtered_out_num_());
+        auto bitset = bitset_;
+        if (bitset.count() == 0) {
+            bitset.set_count(bitset.get_filtered_out_num_());
+        }
 
 #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT)
         const BaseConfig& b_cfg = static_cast<const BaseConfig&>(*cfg);