Method SearchWithRequest

src/algorithm/brute_force.cpp:193–274

}

DatasetPtr
BruteForce::SearchWithRequest(const SearchRequest& request) const {
    std::shared_lock read_lock(this->global_mutex_);

    auto computer = this->inner_codes_->FactoryComputer(request.query_->GetFloat32Vectors());
    DistHeapPtr heap = nullptr;
    ExecutorPtr executor = nullptr;
    Filter* attr_filter = nullptr;
    Filter* filter = nullptr;
    if (request.filter_ != nullptr) {
        filter = request.filter_.get();
    }
    if (request.enable_attribute_filter_) {
        auto& schema = this->attr_filter_index_->field_type_map_;
        auto expr = AstParse(request.attribute_filter_str_, &schema);
        executor = Executor::MakeInstance(this->allocator_, expr, this->attr_filter_index_);
        executor->Init();
        executor->Clear();
        attr_filter = executor->Run();
    }

    std::atomic<uint32_t> dist_cmp{0};

    auto brute_force_params = BruteForceSearchParameters::FromJson(request.params_str_);
    auto parallel_count = brute_force_params.parallel_search_thread_count;
    std::vector<DistHeapPtr> heaps(parallel_count);
    for (auto& cur_heap : heaps) {
        cur_heap = DistanceHeap::MakeInstanceBySize<true, true>(this->allocator_, request.topk_);
    }
    auto search_func = [&](InnerIdType start, InnerIdType end, const DistHeapPtr& cur_heap) {
        float cur_min_dist = std::numeric_limits<float>::max();
        uint32_t dist_cmp_local = 0;
        for (InnerIdType i = start; i < end; ++i) {
            float dist = 0.0F;
            if (attr_filter != nullptr and not attr_filter->CheckValid(i)) {
                continue;
            }
            if (filter == nullptr or filter->CheckValid(this->label_table_->GetLabelById(i))) {
                inner_codes_->Query(&dist, computer, &i, 1);
                ++dist_cmp_local;
                cur_heap->Push(dist, i);
            }
        }

        dist_cmp.fetch_add(dist_cmp_local, std::memory_order_relaxed);
    };

    if (parallel_count == 1 || this->thread_pool_ == nullptr) {
        search_func(0, total_count_, heaps[0]);
        heap = heaps[0];
    } else {
        std::vector<std::future<void>> futures;
        auto chunk_size = (total_count_ + parallel_count - 1) / parallel_count;
        for (auto i = 0; i < parallel_count; ++i) {
            auto start = i * chunk_size;
            auto end = std::min(start + chunk_size, total_count_);
            auto future = this->thread_pool_->GeneralEnqueue(search_func, start, end, heaps[i]);
            futures.emplace_back(std::move(future));