Method search_impl

src/algorithm/pyramid.cpp:317–412

}

DatasetPtr
Pyramid::search_impl(const DatasetPtr& query,
                     const SearchFunc& search_func,
                     InnerSearchParam& search_param) const {
    SearchStatistics stats;
    QueryContext ctx{.stats = &stats};

    const auto* query_path = query->GetPaths();
    CHECK_ARGUMENT(  // NOLINT
        query_path != nullptr || root_->status_ != IndexNode::Status::NO_INDEX,
        "query_path is required when level0 is not built");
    CHECK_ARGUMENT(query->GetFloat32Vectors() != nullptr, "query vectors is required");

    DistHeapPtr search_result = std::make_shared<StandardHeap<true, false>>(allocator_, -1);

    std::shared_lock<std::shared_mutex> lock(resize_mutex_);
    auto vl = pool_->TakeOne();
    if (query_path != nullptr) {
        std::vector<std::future<void>> futures;
        const std::string& current_path = query_path[0];
        auto parsed_path = parse_path(current_path);
        Vector<DistHeapPtr> search_result_lists(parsed_path.size(), allocator_);
        for (uint32_t i = 0; i < parsed_path.size(); ++i) {
            const auto& one_path = parsed_path[i];
            search_result_lists[i] = std::make_shared<StandardHeap<true, false>>(allocator_, -1);
            std::shared_ptr<IndexNode> node = root_;
            bool valid = true;
            for (const auto& item : one_path) {
                node = node->GetChild(item, false);
                if (node == nullptr) {
                    valid = false;
                    break;
                }
            }
            if (valid) {
                if (thread_pool_ != nullptr && search_param.parallel_search_thread_count > 1) {
                    futures.push_back(thread_pool_->GeneralEnqueue([&, node, i]() -> void {
                        node->Search(search_func, vl, search_result_lists[i], search_param.ef);
                    }));
                } else {
                    node->Search(search_func, vl, search_result_lists[i], search_param.ef);
                }
            }
        }

        for (auto& future : futures) {
            future.get();
        }

        for (uint32_t i = 0; i < search_result_lists.size(); ++i) {
            if (i != 0) {
                search_result->Merge(*search_result_lists[i]);
            } else {
                search_result = search_result_lists[i];
            }
        }

    } else {