From a73f22ee5312176a88cdf95f2c09a9ba545b12f9 Mon Sep 17 00:00:00 2001 From: Matthijs Douze Date: Mon, 2 Feb 2026 11:16:50 -0800 Subject: [PATCH] Demo for a diversity filter (#4765) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4765 There are use cases where there is a group id associated to each vector id. We want no more than max_per_group results per group but of course we also want to keep the best results for each group. This diff demonstrates how to obtain this functionality, using a custom result handler. Differential Revision: D91020690 fbshipit-source-id: 92669cf8396ad8b11faa87c4420b35a3db2de7f9 --- .github/workflows/build-pull-request.yml | 3 + demos/CMakeLists.txt | 2 + demos/diversity_filter/CMakeLists.txt | 8 + .../demo_diversity_result_handler.cpp | 298 ++++++++++++++++++ .../diversity_result_handler.h | 116 +++++++ 5 files changed, 427 insertions(+) create mode 100644 demos/diversity_filter/CMakeLists.txt create mode 100644 demos/diversity_filter/demo_diversity_result_handler.cpp create mode 100644 demos/diversity_filter/diversity_result_handler.h diff --git a/.github/workflows/build-pull-request.yml b/.github/workflows/build-pull-request.yml index db01fcdb0..d10ad3de2 100644 --- a/.github/workflows/build-pull-request.yml +++ b/.github/workflows/build-pull-request.yml @@ -38,6 +38,9 @@ jobs: uses: actions/checkout@v4 - name: Build and Test (cmake) uses: ./.github/actions/build_cmake + - name: Build C++ demos + run: | + make -C build demo_diversity_result_handler linux-x86_64-AVX2-cmake: name: Linux x86_64 AVX2 (cmake) needs: linux-x86_64-cmake diff --git a/demos/CMakeLists.txt b/demos/CMakeLists.txt index ec9de3161..31445d535 100644 --- a/demos/CMakeLists.txt +++ b/demos/CMakeLists.txt @@ -23,3 +23,5 @@ target_link_libraries(demo_weighted_kmeans PRIVATE faiss) add_executable(demo_residual_quantizer EXCLUDE_FROM_ALL demo_residual_quantizer.cpp) target_link_libraries(demo_residual_quantizer PRIVATE faiss) + +add_subdirectory(diversity_filter) diff --git a/demos/diversity_filter/CMakeLists.txt b/demos/diversity_filter/CMakeLists.txt new file mode 100644 index 000000000..9aa8502b3 --- /dev/null +++ b/demos/diversity_filter/CMakeLists.txt @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +add_executable(demo_diversity_result_handler EXCLUDE_FROM_ALL + demo_diversity_result_handler.cpp) +target_link_libraries(demo_diversity_result_handler PRIVATE faiss) diff --git a/demos/diversity_filter/demo_diversity_result_handler.cpp b/demos/diversity_filter/demo_diversity_result_handler.cpp new file mode 100644 index 000000000..91e36c55d --- /dev/null +++ b/demos/diversity_filter/demo_diversity_result_handler.cpp @@ -0,0 +1,298 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include + +#include +#include + +#include "diversity_result_handler.h" + +// compare results, taking into account that ties may be ordered arbitrarily +bool compare_results( + int test_no, + int K, + const std::vector& reference_D, + const std::vector& reference_I, + const std::vector& new_D, + const std::vector& new_I, + const std::vector& group_ids) { + float current_dis = NAN; + std::set ref_ids; + std::set new_ids; + + for (int i = 0; i < K; i++) { + float d = reference_D[i]; + if (new_D[i] != d) { + printf("test_no=%d Distance mismatch at %d: %g vs %g\n", + test_no, + i, + new_D[i], + reference_D[i]); + return false; + } + if (d != current_dis) { + // check that the ids for the previous distance are the same + if (ref_ids != new_ids) { + // find the extra elements from each set + std::set extra_ref_ids; + std::set extra_new_ids; + std::set_difference( + ref_ids.begin(), + ref_ids.end(), + new_ids.begin(), + new_ids.end(), + std::inserter(extra_ref_ids, extra_ref_ids.end())); + std::set_difference( + new_ids.begin(), + new_ids.end(), + ref_ids.begin(), + ref_ids.end(), + std::inserter(extra_new_ids, extra_new_ids.end())); + if (extra_ref_ids.size() == 1 && extra_new_ids.size() == 1) { + int64_t id1 = *extra_ref_ids.begin(); + int64_t id2 = *extra_new_ids.begin(); + if (group_ids[id1] == group_ids[id2]) { + // this is a valid case - same group, just different + // tie-breaking + goto ok; + } + } + + printf("test_no=%d i=%d Id mismatch at distance %g\n", + test_no, + i, + current_dis); + printf(" ref_ids: {"); + for (auto id : ref_ids) { + printf(" %ld", id); + } + printf(" }\n"); + printf(" new_ids: {"); + for (auto id : new_ids) { + printf(" %ld", id); + } + printf(" }\n"); + printf("\n=== Reference results (all K) ===\n"); + for (int j = 0; j < K; j++) { + printf(" [%d] id=%zu distance=%g group=%d\n", + j, + reference_I[j], + reference_D[j], + group_ids[reference_I[j]]); + } + printf("\n=== New handler results (all K) ===\n"); + for (int j = 0; j < K; j++) { + printf(" [%d] id=%ld distance=%g group=%d\n", + j, + new_I[j], + new_D[j], + group_ids[new_I[j]]); + } + return false; + } + ok: + current_dis = d; + ref_ids.clear(); + new_ids.clear(); + } + + ref_ids.insert(reference_I[i]); + new_ids.insert(new_I[i]); + } + return true; +} + +// test result handler with a table of random distances and group ids, without +// reference to an index +void test_result_handler_table() { + int N = 100; + int K = 10; + int n_groups = 5; + int max_per_group = 3; + + // Fixed seed for reproducibility + srand(12345); + + for (int test_no = 0; test_no < 200; test_no++) { + std::vector group_ids(N); + std::vector distance(N); + + for (int i = 0; i < N; i++) { + group_ids[i] = rand() % n_groups; + // distance[i] = (float)rand() / (float)RAND_MAX; + // increase likelihood of same-distances + distance[i] = rand() % 500 / 500.0; + } + + // reference implementation + std::vector reference_I(K); + std::vector reference_D(K); + { + std::vector permutation(N); + std::iota(permutation.begin(), permutation.end(), 0); + std::sort( + permutation.begin(), + permutation.end(), + [&](size_t i, size_t j) { + return distance[i] < distance[j]; + }); + + std::unordered_map group_sizes; + size_t n_res = 0; + size_t perm_idx = 0; + while (n_res < K && perm_idx < N) { + size_t idx = permutation[perm_idx]; + int group_id = group_ids[idx]; + perm_idx++; + if (group_sizes[group_id] < max_per_group) { + group_sizes[group_id]++; + reference_I[n_res] = idx; + reference_D[n_res] = distance[idx]; + n_res++; + } + } + } + + // run result_handler + std::vector new_I(K); + std::vector new_D(K); + { + DiversityResultHandlerBubble handler( + group_ids, max_per_group, K, new_D.data(), new_I.data()); + for (int i = 0; i < N; i++) { + handler.add_result(distance[i], i); + } + } + + if (!compare_results( + test_no, + K, + reference_D, + reference_I, + new_D, + new_I, + group_ids)) { + printf("test_no=%d FAILED\n", test_no); + return; + } + } + + printf("All tests passed!\n"); +} + +void test_diversity_index(const char* factory_string) { + int d = 64; // dimension + int N = 2000; // database size + int nt = 5000; // training set size + int nq = 100; // number of queries + int K = 10; // number of results to return + int n_groups = 5; + int max_per_group = 3; + + printf("Testing diversity index with factory: %s\n", factory_string); + + // Create index using factory + std::unique_ptr index( + faiss::index_factory(d, factory_string, faiss::METRIC_L2)); + + // Generate training data + std::vector xt(nt * d); + faiss::rand_smooth_vectors(nt, d, xt.data(), 1234); + + // Generate database vectors + std::vector xb(N * d); + faiss::rand_smooth_vectors(N, d, xb.data(), 4567); + + // Generate query vectors + std::vector xq(nq * d); + faiss::rand_smooth_vectors(nq, d, xq.data(), 7890); + + // Train the index + index->train(nt, xt.data()); + + // Add database vectors + index->add(N, xb.data()); + + // Set nprobe for IVF indexes + if (auto* ivf = dynamic_cast(index.get())) { + ivf->nprobe = 4; + } + + // Generate group ids for each database vector + srand(54321); + std::vector group_ids(N); + for (int i = 0; i < N; i++) { + group_ids[i] = rand() % n_groups; + } + + // Test each query + for (int q = 0; q < nq; q++) { + const float* query = xq.data() + q * d; + + // Reference: first get all results, then filter for diversity + std::vector all_D(N); + std::vector all_I(N); + + // Get all distances using a standard search with large k + // We'll use the index's search to get ordered results + int large_k = std::min(N, 500); // use a reasonably large k + std::vector ref_D(large_k); + std::vector ref_I(large_k); + index->search(1, query, large_k, ref_D.data(), ref_I.data()); + + // Reference implementation: filter for diversity + std::vector reference_I(K); + std::vector reference_D(K); + { + std::unordered_map group_sizes; + size_t n_res = 0; + for (int i = 0; i < large_k && n_res < K; i++) { + faiss::idx_t idx = ref_I[i]; + if (idx < 0) { + continue; + } + int group_id = group_ids[idx]; + if (group_sizes[group_id] < max_per_group) { + group_sizes[group_id]++; + reference_I[n_res] = idx; + reference_D[n_res] = ref_D[i]; + n_res++; + } + } + } + + // Now use the DiversityResultHandlerBubble with search1 + std::vector new_D(K); + std::vector new_I(K); + { + DiversityResultHandlerBubble handler( + group_ids, max_per_group, K, new_D.data(), new_I.data()); + handler.threshold = std::numeric_limits::max(); + index->search1(query, handler); + } + + // Compare results + if (!compare_results( + q, K, reference_D, reference_I, new_D, new_I, group_ids)) { + printf("test_diversity_index: query %d FAILED\n", q); + return; + } + } + + printf("test_diversity_index: All %d queries passed!\n", nq); +} + +int main() { + test_result_handler_table(); + test_diversity_index("IVF32,Flat"); + return 0; +} diff --git a/demos/diversity_filter/diversity_result_handler.h b/demos/diversity_filter/diversity_result_handler.h new file mode 100644 index 000000000..eca28975e --- /dev/null +++ b/demos/diversity_filter/diversity_result_handler.h @@ -0,0 +1,116 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include + +using TheC = faiss::CMax; + +/** There is a group id associated to each vector id. We want no more than + * max_per_group results per group but of course we also want to keep the best + * results for each group. + * This implementation is a stateless handler that always keeps results sorted, + * which means updates are O(k) */ +struct DiversityResultHandlerBubble : faiss::ResultHandlerT { + using group_t = int; + + // records the group for each database vector + const std::vector& id_to_group; + + // we want no more than this many results per group + size_t max_per_group; + + // the distances and labels of the results + size_t K; // number of results to keep + float* distances; + int64_t* labels; + + // found so far + size_t nresults = 0; + + DiversityResultHandlerBubble( + const std::vector& id_to_group, + size_t max_per_group, + size_t K, + float* distances, + int64_t* labels) + : id_to_group(id_to_group), + max_per_group(max_per_group), + K(K), + distances(distances), + labels(labels) {} + + /// find where the element distance should be inserted in the sorted array + /// distances + size_t bissection(float distance) { + size_t lo = 0, hi = nresults; + while (lo < hi) { + size_t mid = (lo + hi) / 2; + if (distances[mid] < distance) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + } + + /// shift results one step to the right in [i0, i1), ie. i1 is overwritten + void shift_results(size_t i0, size_t i1) { + if (i0 >= i1) { + return; + } + memmove(&distances[i0 + 1], &distances[i0], (i1 - i0) * sizeof(float)); + memmove(&labels[i0 + 1], &labels[i0], (i1 - i0) * sizeof(int64_t)); + } + + virtual bool add_result(float distance, faiss::idx_t i) final { + group_t group = id_to_group[i]; + if (nresults == K && distance >= distances[K - 1]) { + return false; + } + // count number of occurrences of this group and keep track of the worst + // result of this group + int n = 0; + size_t worst_i = 0; + for (int j = 0; j < nresults; j++) { + if (id_to_group[labels[j]] == group) { + n++; + worst_i = j; // is the worst since sorted + } + } + if (n >= max_per_group && distance >= distances[worst_i]) { + return false; + } + // we can add the new result + size_t insertion_point = bissection(distance); + if (n >= max_per_group) { + // replace worst_i with new result + // shift elements between insertion_point and worst_i + if (insertion_point < worst_i) { + shift_results(insertion_point, worst_i); + } + } else if (nresults < (int)K) { + // add new result + shift_results(insertion_point, nresults); + nresults++; + } else { + // replace the last result + shift_results(insertion_point, K - 1); + } + distances[insertion_point] = distance; + labels[insertion_point] = i; + return true; + } +}; + +/** TODO implement a version with a double heap: the first heap keeps track of + * the worst result of each group, K other heaps keep track of the results of + * each group. This would ensure log(k) update time. */