std::vector<float> D;

std::vector<idx_t> I;

bool add_result(float distance, idx_t i) override {

// we never change the default threshold so that all results pass

D.push_back(distance);

I.push_back(i);

return true;

}

// Sort results by distance and return the top k

void get_top_k(

size_t top_k,

std::vector<float>& out_D,

std::vector<idx_t>& out_I,

bool is_max) const {

size_t n_results = D.size();

// Create indices for sorting

std::vector<size_t> indices(n_results);

for (size_t i = 0; i < n_results; i++) {

indices[i] = i;

}

// Make local copies for lambda capture

const std::vector<float>& distances = D;

// Sort by distance (ascending for L2, descending for IP)

if (is_max) {

std::sort(

indices.begin(),

indices.end(),

[&distances](size_t a, size_t b) {

return distances[a] > distances[b];

});

} else {

std::sort(

indices.begin(),

indices.end(),

[&distances](size_t a, size_t b) {

return distances[a] < distances[b];

});

}

// Take top k

size_t n_out = std::min(top_k, n_results);

out_D.resize(n_out);

out_I.resize(n_out);

for (size_t i = 0; i < n_out; i++) {

out_D[i] = D[indices[i]];

out_I[i] = I[indices[i]];

}

}

std::vector<float> data(n * d);

rand_smooth_vectors(n, d, data.data(), seed);

return data;

const char* index_key,

MetricType metric,

double min_match_ratio = 1.0) {

// Create index using factory

std::unique_ptr<Index> index(index_factory(d, index_key, metric));

ASSERT_NE(index, nullptr) << "Failed to create IVF index for " << index_key;

// Generate smooth random data for training and database

auto xt = make_smooth_data(nt, 1234);

auto xb = make_smooth_data(nb, 4567);

auto xq = make_smooth_data(nq, 7890);

// Train the index

index->train(nt, xt.data());

// Add database vectors

index->add(nb, xb.data());

// Set nprobe for IVF indexes

if (IndexIVF* index_ivf = dynamic_cast<IndexIVF*>(index.get())) {

index_ivf->nprobe = 4;

}

// Perform reference search using standard search method

std::vector<idx_t> Iref(nq * k);

std::vector<float> Dref(nq * k);

index->search(nq, xq.data(), k, Dref.data(), Iref.data());

// For IP metric, we need to sort in descending order

bool is_max = (metric == METRIC_INNER_PRODUCT);

// Now test search1 with custom handler for each query

for (size_t q = 0; q < nq; q++) {

CollectAllResultHandler handler;

// Set threshold to collect all results

// For L2: use max float (condition: threshold > distance)

// For IP: use min float (condition: threshold < distance)

if (is_max) {

handler.threshold = -std::numeric_limits<float>::max();

} else {

handler.threshold = std::numeric_limits<float>::max();

}

index->search1(xq.data() + q * d, handler);

// Sort the handler results and get top k

std::vector<float> D_handler;

std::vector<idx_t> I_handler;

handler.get_top_k(k, D_handler, I_handler, is_max);

// Compare with reference results for this query using set comparison

const idx_t* Iref_q = Iref.data() + q * k;

// Create sets of IDs for comparison

std::unordered_set<idx_t> ref_set(Iref_q, Iref_q + k);

std::unordered_set<idx_t> handler_set(

I_handler.begin(), I_handler.end());

// Count how many results from handler are in the reference set

int matches = 0;

for (idx_t id : handler_set) {

if (ref_set.count(id) > 0) {

matches++;

}

}

// Check that matches meet the minimum threshold

EXPECT_GE(matches, static_cast<int>(k * min_match_ratio))

<< "Query " << q << ": expected at least "

<< static_cast<int>(k * min_match_ratio) << " matches, got "

<< matches;

}