Support MV only for HNSW

Signed-off-by: chasingegg <[email protected]>
zilliztech · Jan 14, 2025 · 6dec4cf · 6dec4cf
1 parent 59bfdc2
commit 6dec4cf
Show file tree

Hide file tree

Showing 13 changed files with 2,306 additions and 1,156 deletions.
diff --git a/include/knowhere/bitsetview.h b/include/knowhere/bitsetview.h
@@ -95,6 +95,38 @@ class BitsetView {
         return ret;
     }
 
+    size_t
+    get_first_valid_index() const {
+        size_t ret = 0;
+        auto len_uint8 = byte_size();
+        auto len_uint64 = len_uint8 >> 3;
+
+        uint64_t* p_uint64 = (uint64_t*)bits_;
+        for (size_t i = 0; i < len_uint64; i++) {
+            uint64_t value = (~(*p_uint64));
+            if (value == 0) {
+                p_uint64++;
+                continue;
+            }
+            ret = __builtin_ctzll(value);
+            return i * 64 + ret;
+        }
+
+        // calculate remainder
+        uint8_t* p_uint8 = (uint8_t*)bits_ + (len_uint64 << 3);
+        for (size_t i = 0; i < len_uint8 - (len_uint64 << 3); i++) {
+            uint8_t value = (~(*p_uint8));
+            if (value == 0) {
+                p_uint8++;
+                continue;
+            }
+            ret = __builtin_ctz(value);
+            return len_uint64 * 64 + i * 8 + ret;
+        }
+
+        return num_bits_;
+    }
+
     std::string
     to_string(size_t from, size_t to) const {
         if (empty()) {

diff --git a/include/knowhere/bitsetview_idselector.h b/include/knowhere/bitsetview_idselector.h
@@ -32,4 +32,21 @@ struct BitsetViewIDSelector final : faiss::IDSelector {
     }
 };
 
+struct BitsetViewWithMappingIDSelector final : faiss::IDSelector {
+    const BitsetView bitset_view;
+    const uint32_t* out_id_mapping;
+    const size_t id_offset;
+
+    inline BitsetViewWithMappingIDSelector(BitsetView bitset_view, const uint32_t* out_id_mapping,
+                                           const size_t offset = 0)
+        : bitset_view{bitset_view}, out_id_mapping(out_id_mapping), id_offset(offset) {
+    }
+
+    inline bool
+    is_member(faiss::idx_t id) const override final {
+        // it is by design that out_id_mapping == nullptr is not tested here
+        return (!bitset_view.test(out_id_mapping[id + id_offset]));
+    }
+};
+
 }  // namespace knowhere
diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h
@@ -95,6 +95,7 @@ constexpr const char* JSON_ID_SET = "json_id_set";
 constexpr const char* TRACE_ID = "trace_id";
 constexpr const char* SPAN_ID = "span_id";
 constexpr const char* TRACE_FLAGS = "trace_flags";
+constexpr const char* SCALAR_INFO = "scalar_info";
 constexpr const char* MATERIALIZED_VIEW_SEARCH_INFO = "materialized_view_search_info";
 constexpr const char* MATERIALIZED_VIEW_OPT_FIELDS_PATH = "opt_fields_path";
 constexpr const char* MAX_EMPTY_RESULT_BUCKETS = "max_empty_result_buckets";