feat: reuse prepared search configurations

Signed-off-by: xianliang.li <xianliang.li@zilliz.com>

Author: foxspy

include/knowhere/index/index_node.h

@@ -12,7 +12,9 @@
 #ifndef INDEX_NODE_H
 #define INDEX_NODE_H
 
+#include <atomic>
 #include <functional>
+#include <memory>
 #include <mutex>
 #include <queue>
 #include <utility>
@@ -618,6 +620,22 @@ class IndexNode : public Object {
     SearchEmbList(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset,
                   milvus::OpContext* op_context = nullptr) const;
 
+ public:
+    virtual bool
+    SupportsSearchConfigCache() const {
+        return false;
+    }
+
+    virtual expected<DataSetPtr>
+    SearchWithPreparedConfig(const DataSetPtr, std::shared_ptr<const Config>, const BitsetView&,
+                             milvus::OpContext* = nullptr) const {
+        return expected<DataSetPtr>::Err(Status::not_implemented, "prepared search config is not supported");
+    }
+
+    expected<std::shared_ptr<const Config>>
+    GetOrCreateSearchConfig(const Json& json) const;
+
+ protected:
     static EmbListMetaHeader
     ParseEmbListMetaHeader(const uint8_t* data, int64_t size);
 
@@ -637,6 +655,17 @@ class IndexNode : public Object {
                        std::shared_ptr<ThreadPool> pool, milvus::OpContext* op_context = nullptr) const;
 
     Version version_;
+
+ private:
+    struct SearchConfigCacheEntry {
+        Json json;
+        std::shared_ptr<const Config> config;
+    };
+
+    mutable std::atomic<std::shared_ptr<const SearchConfigCacheEntry>> search_config_cache_;
+    mutable std::mutex search_config_cache_mutex_;
+
+ protected:
     std::shared_ptr<EmbListOffset> emb_list_offset_;  // emb_list group offset structure (shared with strategy)
     std::string el_metric_type_;
     EmbListStrategyPtr emb_list_strategy_;  // emb_list encoding strategy (tokenann/muvera)

src/index/index.cc

@@ -133,12 +133,23 @@ inline expected<DataSetPtr>
 Index<T>::Search(const DataSetPtr dataset, const Json& json, const BitsetView& bitset_,
                  milvus::OpContext* op_context) const noexcept {
     return GuardedCall([&]() -> expected<DataSetPtr> {
-        auto cfg = this->node->CreateConfig();
+        std::unique_ptr<BaseConfig> owned_cfg;
+        std::shared_ptr<const Config> prepared_cfg;
         std::string msg;
-        const Status load_status = LoadConfig(cfg.get(), json, knowhere::SEARCH, "Search", &msg);
-        if (load_status != Status::success) {
-            return expected<DataSetPtr>::Err(load_status, msg);
+        if (this->node->SupportsSearchConfigCache()) {
+            auto result = this->node->GetOrCreateSearchConfig(json);
+            if (!result.has_value()) {
+                return expected<DataSetPtr>::Err(result.error(), result.what());
+            }
+            prepared_cfg = std::move(result.value());
+        } else {
+            owned_cfg = this->node->CreateConfig();
+            const Status load_status = LoadConfig(owned_cfg.get(), json, knowhere::SEARCH, "Search", &msg);
+            if (load_status != Status::success) {
+                return expected<DataSetPtr>::Err(load_status, msg);
+            }
         }
+        const Config* cfg = prepared_cfg != nullptr ? prepared_cfg.get() : owned_cfg.get();
         // when index is immutable, bitset size should always equal to data count in index
         // when index is mutable, it could happen that data count larger than bitset size, see
         // https://github.com/zilliztech/knowhere/issues/70
@@ -177,14 +188,18 @@ Index<T>::Search(const DataSetPtr dataset, const Json& json, const BitsetView& b
         // LCOV_EXCL_STOP
 
         TimeRecorder rc("Search");
-        auto k = cfg->k.value();
-        auto res = this->node->SearchEmbListIfNeed(dataset, std::move(cfg), bitset, op_context);
+        auto k = b_cfg.k.value();
+        auto res = prepared_cfg != nullptr
+                       ? this->node->SearchWithPreparedConfig(dataset, std::move(prepared_cfg), bitset, op_context)
+                       : this->node->SearchEmbListIfNeed(dataset, std::move(owned_cfg), bitset, op_context);
         auto time = rc.ElapseFromBegin("done");
         time *= 0.001;  // convert to ms
         this->node->GetSearchLatencyMetric().Observe(time);
         knowhere_search_topk.Observe(k);
 #else
-        auto res = this->node->SearchEmbListIfNeed(dataset, std::move(cfg), bitset, op_context);
+        auto res = prepared_cfg != nullptr
+                       ? this->node->SearchWithPreparedConfig(dataset, std::move(prepared_cfg), bitset, op_context)
+                       : this->node->SearchEmbListIfNeed(dataset, std::move(owned_cfg), bitset, op_context);
 #endif
         return res;
     });

src/index/index_node.cc

@@ -32,6 +32,39 @@
 
 namespace knowhere {
 
+expected<std::shared_ptr<const Config>>
+IndexNode::GetOrCreateSearchConfig(const Json& json) const {
+    auto cached = search_config_cache_.load(std::memory_order_acquire);
+    if (cached != nullptr && cached->json == json) {
+        return cached->config;
+    }
+
+    std::lock_guard lock(search_config_cache_mutex_);
+    cached = search_config_cache_.load(std::memory_order_relaxed);
+    if (cached != nullptr && cached->json == json) {
+        return cached->config;
+    }
+
+    auto cfg = CreateConfig();
+    Json normalized_json(json);
+    std::string msg;
+    auto status = Config::FormatAndCheck(*cfg, normalized_json, &msg);
+    LOG_KNOWHERE_DEBUG_ << "Search config dump: " << normalized_json.dump();
+    if (status != Status::success) {
+        return expected<std::shared_ptr<const Config>>::Err(status, msg);
+    }
+    cfg->CaptureRawJson(normalized_json);
+    status = Config::Load(*cfg, normalized_json, knowhere::SEARCH, &msg);
+    if (status != Status::success) {
+        return expected<std::shared_ptr<const Config>>::Err(status, msg);
+    }
+
+    std::shared_ptr<const Config> prepared_config(std::move(cfg));
+    auto entry = std::make_shared<const SearchConfigCacheEntry>(SearchConfigCacheEntry{json, prepared_config});
+    search_config_cache_.store(std::move(entry), std::memory_order_release);
+    return prepared_config;
+}
+
 // NOLINTBEGIN(google-default-arguments)
 expected<DataSetPtr>
 IndexNode::RangeSearch(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset,

tests/ut/test_index_node.cc

@@ -9,6 +9,8 @@
 // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 // or implied. See the License for the specific language governing permissions and limitations under the License.
 
+#include <atomic>
+#include <future>
 #include <unordered_set>
 
 #include "catch2/catch_approx.hpp"
@@ -136,6 +138,49 @@ class BaseFlatIndexNode : public IndexNode {
     }
 };
 
+template <typename DataType>
+class CachedSearchConfigIndexNode : public BaseFlatIndexNode<DataType> {
+ public:
+    CachedSearchConfigIndexNode(const int32_t& version, const Object& object)
+        : BaseFlatIndexNode<DataType>(version, object) {
+    }
+
+    bool
+    SupportsSearchConfigCache() const override {
+        return true;
+    }
+
+    expected<DataSetPtr>
+    SearchWithPreparedConfig(const DataSetPtr, std::shared_ptr<const Config> cfg, const BitsetView&,
+                             milvus::OpContext*) const override {
+        const Config* expected = nullptr;
+        first_config_.compare_exchange_strong(expected, cfg.get());
+        reused_same_config_.store(reused_same_config_.load() && first_config_.load() == cfg.get());
+        return std::make_shared<DataSet>();
+    }
+
+    std::unique_ptr<BaseConfig>
+    CreateConfig() const override {
+        create_config_calls_.fetch_add(1);
+        return std::make_unique<BaseConfig>();
+    }
+
+    int
+    CreateConfigCalls() const {
+        return create_config_calls_.load();
+    }
+
+    bool
+    ReusedSameConfig() const {
+        return reused_same_config_.load();
+    }
+
+ private:
+    mutable std::atomic<int> create_config_calls_{0};
+    mutable std::atomic<const Config*> first_config_{nullptr};
+    mutable std::atomic<bool> reused_same_config_{true};
+};
+
 TEST_CASE("Test index node") {
     auto version = GenTestVersionList();
     DataSetPtr ds = std::make_shared<DataSet>();
@@ -208,3 +253,44 @@ TEST_CASE("Test index node") {
     }
 #pragma GCC diagnostic pop
 }
+
+TEST_CASE("Search reuses an immutable prepared config", "[search_config_cache]") {
+    KNOWHERE_SIMPLE_REGISTER_GLOBAL(SEARCH_CONFIG_CACHE, CachedSearchConfigIndexNode, fp32, knowhere::feature::FLOAT32);
+    const auto version = GenTestVersionList();
+    auto dataset = std::make_shared<DataSet>();
+    const Json base_search_config = {{meta::METRIC_TYPE, metric::L2}, {meta::TOPK, 10}};
+
+    SECTION("same config reuses the prepared object") {
+        auto index = IndexFactory::Instance().Create<fp32>("SEARCH_CONFIG_CACHE", version).value();
+        auto* node = dynamic_cast<CachedSearchConfigIndexNode<fp32>*>(index.Node());
+        REQUIRE(node != nullptr);
+
+        REQUIRE(index.Search(dataset, base_search_config, nullptr).has_value());
+        REQUIRE(index.Search(dataset, base_search_config, nullptr).has_value());
+        REQUIRE(node->CreateConfigCalls() == 1);
+        REQUIRE(node->ReusedSameConfig());
+
+        auto changed_search_config = base_search_config;
+        changed_search_config[meta::TOPK] = 20;
+        REQUIRE(index.Search(dataset, changed_search_config, nullptr).has_value());
+        REQUIRE(node->CreateConfigCalls() == 2);
+    }
+
+    SECTION("concurrent searches prepare the config once") {
+        auto index = IndexFactory::Instance().Create<fp32>("SEARCH_CONFIG_CACHE", version).value();
+        auto* node = dynamic_cast<CachedSearchConfigIndexNode<fp32>*>(index.Node());
+        REQUIRE(node != nullptr);
+
+        std::vector<std::future<expected<DataSetPtr>>> searches;
+        for (int i = 0; i < 32; ++i) {
+            searches.emplace_back(
+                std::async(std::launch::async, [&] { return index.Search(dataset, base_search_config, nullptr); }));
+        }
+        for (auto& search : searches) {
+            REQUIRE(search.get().has_value());
+        }
+
+        REQUIRE(node->CreateConfigCalls() == 1);
+        REQUIRE(node->ReusedSameConfig());
+    }
+}