diff --git a/benchs/bench_hnsw_flat_panorama.py b/benchs/bench_hnsw_flat_panorama.py new file mode 100644 index 000000000..ea6710ce0 --- /dev/null +++ b/benchs/bench_hnsw_flat_panorama.py @@ -0,0 +1,153 @@ +# 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. + +import multiprocessing as mp +import time + +import faiss +import matplotlib.pyplot as plt +import numpy as np + +try: + from faiss.contrib.datasets_fb import ( + DatasetSIFT1M, + DatasetGIST1M, + SyntheticDataset, + ) +except ImportError: + from faiss.contrib.datasets import ( + DatasetSIFT1M, + DatasetGIST1M, + SyntheticDataset, + ) + + +def eval_recall(index, efSearch_val, xq, gt, k): + """Evaluate recall and QPS for a given efSearch value.""" + t0 = time.time() + _, I = index.search(xq, k=k) + t = time.time() - t0 + speed = t * 1000 / len(xq) + qps = 1000 / speed + + corrects = (gt == I).sum() + recall = corrects / (len(xq) * k) + print( + f"\tefSearch {efSearch_val:3d}, Recall@{k}: " + f"{recall:.6f}, speed: {speed:.6f} ms/query, QPS: {qps:.2f}" + ) + + return recall, qps + + +def get_hnsw_index(index): + """Extract the underlying HNSW index from a PreTransform index.""" + if isinstance(index, faiss.IndexPreTransform): + return faiss.downcast_index(index.index) + return index + + +def eval_and_plot(name, ds, k=10, nlevels=8, plot_data=None): + """Evaluate an index configuration and collect data for plotting.""" + xq = ds.get_queries() + xb = ds.get_database() + gt = ds.get_groundtruth() + + if hasattr(ds, "get_train"): + xt = ds.get_train() + else: + # Use database as training data if no separate train set + xt = xb + + nb, d = xb.shape + nq, d = xq.shape + gt = gt[:, :k] + + print(f"\n======{name} on {ds.__class__.__name__}======") + print(f"Database: {nb} vectors, {d} dimensions") + print(f"Queries: {nq} vectors") + + # Create index + index = faiss.index_factory(d, name) + + faiss.omp_set_num_threads(mp.cpu_count()) + index.train(xt) + index.add(xb) + + faiss.omp_set_num_threads(1) + + # Get the underlying HNSW index for setting efSearch + hnsw_index = get_hnsw_index(index) + + data = [] + for efSearch in [16, 32, 64, 128, 256, 512]: + hnsw_index.hnsw.efSearch = efSearch + recall, qps = eval_recall(index, efSearch, xq, gt, k) + data.append((recall, qps)) + + if plot_data is not None: + data = np.array(data) + plot_data.append((name, data)) + + +def benchmark_dataset(ds, dataset_name, k=10, nlevels=8, M=32): + """Benchmark both regular HNSW and HNSW Panorama on a dataset.""" + d = ds.d + + plot_data = [] + + # HNSW Flat (baseline) + eval_and_plot(f"HNSW{M},Flat", ds, k=k, nlevels=nlevels, plot_data=plot_data) + + # HNSW Flat Panorama (with PCA to concentrate energy) + eval_and_plot( + f"PCA{d},HNSW{M},FlatPanorama{nlevels}", + ds, + k=k, + nlevels=nlevels, + plot_data=plot_data, + ) + + # Plot results + plt.figure(figsize=(8, 6), dpi=80) + for name, data in plot_data: + plt.plot(data[:, 0], data[:, 1], marker="o", label=name) + + plt.title(f"HNSW Indexes on {dataset_name}") + plt.xlabel(f"Recall@{k}") + plt.ylabel("QPS") + plt.yscale("log") + plt.legend(bbox_to_anchor=(1.02, 0.1), loc="upper left", borderaxespad=0) + plt.grid(True, alpha=0.3) + + output_file = f"bench_hnsw_flat_panorama_{dataset_name}.png" + plt.savefig(output_file, bbox_inches="tight") + print(f"Saved plot to {output_file}") + plt.close() + + +if __name__ == "__main__": + k = 10 + nlevels = 8 + M = 32 + + # Test on 3 datasets with varying dimensionality: + # SIFT1M (128d), GIST1M (960d), and Synthetic high-dim (2048d) + datasets = [ + (DatasetSIFT1M(), "SIFT1M"), + (DatasetGIST1M(), "GIST1M"), + # Synthetic high-dimensional dataset: 2048d, 100k train, 1M database, 10k queries + (SyntheticDataset(2048, 100000, 1000000, 10000), "Synthetic2048D"), + ] + + for ds, name in datasets: + print(f"\n{'='*60}") + print(f"Benchmarking on {name}") + print(f"{'='*60}") + benchmark_dataset(ds, name, k=k, nlevels=nlevels, M=M) + + print("\n" + "="*60) + print("All benchmarks completed!") + print("="*60) diff --git a/faiss/IndexAdditiveQuantizer.cpp b/faiss/IndexAdditiveQuantizer.cpp index f9e7c773e..a64d36d4d 100644 --- a/faiss/IndexAdditiveQuantizer.cpp +++ b/faiss/IndexAdditiveQuantizer.cpp @@ -86,6 +86,7 @@ struct AQDistanceComputerLUT : FlatCodesDistanceComputer { float bias; void set_query(const float* x) final { + q = x; // this is quite sub-optimal for multiple queries aq.compute_LUT(1, x, LUT.data()); if (is_IP) { diff --git a/faiss/IndexBinaryHNSW.cpp b/faiss/IndexBinaryHNSW.cpp index dabaa5829..dbb7cc319 100644 --- a/faiss/IndexBinaryHNSW.cpp +++ b/faiss/IndexBinaryHNSW.cpp @@ -227,7 +227,11 @@ void IndexBinaryHNSW::search( for (idx_t i = 0; i < n; i++) { res.begin(i); dis->set_query((float*)(x + i * code_size)); - hnsw.search(*dis, res, vt); + // Given that IndexBinaryHNSW is not an IndexHNSW, we pass nullptr + // as the index parameter. This state does not get used in the + // search function, as it is merely there to to enable Panorama + // execution for IndexHNSWFlatPanorama. + hnsw.search(*dis, nullptr, res, vt); res.end(); } } diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index c2ee8567c..bee8b9d9a 100644 --- a/faiss/IndexFlat.cpp +++ b/faiss/IndexFlat.cpp @@ -103,15 +103,24 @@ namespace { struct FlatL2Dis : FlatCodesDistanceComputer { size_t d; idx_t nb; - const float* q; const float* b; size_t ndis; + size_t npartial_dot_products; float distance_to_code(const uint8_t* code) final { ndis++; return fvec_L2sqr(q, (float*)code, d); } + float partial_dot_product( + const idx_t i, + const uint32_t offset, + const uint32_t num_components) final override { + npartial_dot_products++; + return fvec_inner_product( + q + offset, b + i * d + offset, num_components); + } + float symmetric_dis(idx_t i, idx_t j) override { return fvec_L2sqr(b + j * d, b + i * d, d); } @@ -119,12 +128,13 @@ struct FlatL2Dis : FlatCodesDistanceComputer { explicit FlatL2Dis(const IndexFlat& storage, const float* q = nullptr) : FlatCodesDistanceComputer( storage.codes.data(), - storage.code_size), + storage.code_size, + q), d(storage.d), nb(storage.ntotal), - q(q), b(storage.get_xb()), - ndis(0) {} + ndis(0), + npartial_dot_products(0) {} void set_query(const float* x) override { q = x; @@ -162,6 +172,50 @@ struct FlatL2Dis : FlatCodesDistanceComputer { dis2 = dp2; dis3 = dp3; } + + void partial_dot_product_batch_4( + const idx_t idx0, + const idx_t idx1, + const idx_t idx2, + const idx_t idx3, + float& dp0, + float& dp1, + float& dp2, + float& dp3, + const uint32_t offset, + const uint32_t num_components) final override { + npartial_dot_products += 4; + + // compute first, assign next + const float* __restrict y0 = + reinterpret_cast(codes + idx0 * code_size); + const float* __restrict y1 = + reinterpret_cast(codes + idx1 * code_size); + const float* __restrict y2 = + reinterpret_cast(codes + idx2 * code_size); + const float* __restrict y3 = + reinterpret_cast(codes + idx3 * code_size); + + float dp0_ = 0; + float dp1_ = 0; + float dp2_ = 0; + float dp3_ = 0; + fvec_inner_product_batch_4( + q + offset, + y0 + offset, + y1 + offset, + y2 + offset, + y3 + offset, + num_components, + dp0_, + dp1_, + dp2_, + dp3_); + dp0 = dp0_; + dp1 = dp1_; + dp2 = dp2_; + dp3 = dp3_; + } }; struct FlatIPDis : FlatCodesDistanceComputer { diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index b7c3bfe10..d51307964 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -276,7 +276,7 @@ void hnsw_search( res.begin(i); dis->set_query(x + i * index->d); - HNSWStats stats = hnsw.search(*dis, res, vt, params); + HNSWStats stats = hnsw.search(*dis, index, res, vt, params); n1 += stats.n1; n2 += stats.n2; ndis += stats.ndis; @@ -649,6 +649,95 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric) is_trained = true; } +/************************************************************** + * IndexHNSWFlatPanorama implementation + **************************************************************/ + +void IndexHNSWFlatPanorama::compute_cum_sums( + const float* x, + float* dst_cum_sums, + int d, + int num_panorama_levels, + int panorama_level_width) { + // Iterate backwards through levels, accumulating sum as we go. + // This avoids computing the suffix sum for each vector, which takes + // extra memory. + + float sum = 0.0f; + dst_cum_sums[num_panorama_levels] = 0.0f; + for (int level = num_panorama_levels - 1; level >= 0; level--) { + int start_idx = level * panorama_level_width; + int end_idx = std::min(start_idx + panorama_level_width, d); + for (int j = start_idx; j < end_idx; j++) { + sum += x[j] * x[j]; + } + dst_cum_sums[level] = std::sqrt(sum); + } +} + +IndexHNSWFlatPanorama::IndexHNSWFlatPanorama() + : IndexHNSWFlat(), + cum_sums(), + panorama_level_width(0), + num_panorama_levels(0) {} + +IndexHNSWFlatPanorama::IndexHNSWFlatPanorama( + int d, + int M, + int num_panorama_levels, + MetricType metric) + : IndexHNSWFlat(d, M, metric), + cum_sums(), + panorama_level_width( + (d + num_panorama_levels - 1) / num_panorama_levels), + num_panorama_levels(num_panorama_levels) { + // For now, we only support L2 distance. + // Supporting dot product and cosine distance is a trivial addition + // left for future work. + FAISS_THROW_IF_NOT(metric == METRIC_L2); + + // Enable Panorama search mode. + // This is not ideal, but is still more simple than making a subclass of + // HNSW and overriding the search logic. + hnsw.is_panorama = true; +} + +void IndexHNSWFlatPanorama::add(idx_t n, const float* x) { + idx_t n0 = ntotal; + cum_sums.resize((ntotal + n) * (num_panorama_levels + 1)); + + for (size_t idx = 0; idx < n; idx++) { + const float* vector = x + idx * d; + compute_cum_sums( + vector, + &cum_sums[(n0 + idx) * (num_panorama_levels + 1)], + d, + num_panorama_levels, + panorama_level_width); + } + + IndexHNSWFlat::add(n, x); +} + +void IndexHNSWFlatPanorama::reset() { + cum_sums.clear(); + IndexHNSWFlat::reset(); +} + +void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) { + std::vector new_cum_sums(ntotal * (num_panorama_levels + 1)); + + for (idx_t i = 0; i < ntotal; i++) { + idx_t src = perm[i]; + memcpy(&new_cum_sums[i * (num_panorama_levels + 1)], + &cum_sums[src * (num_panorama_levels + 1)], + (num_panorama_levels + 1) * sizeof(float)); + } + + std::swap(cum_sums, new_cum_sums); + IndexHNSWFlat::permute_entries(perm); +} + /************************************************************** * IndexHNSWPQ implementation **************************************************************/ diff --git a/faiss/IndexHNSW.h b/faiss/IndexHNSW.h index 727210a1b..395d185b1 100644 --- a/faiss/IndexHNSW.h +++ b/faiss/IndexHNSW.h @@ -111,7 +111,7 @@ struct IndexHNSW : Index { void link_singletons(); - void permute_entries(const idx_t* perm); + virtual void permute_entries(const idx_t* perm); DistanceComputer* get_distance_computer() const override; }; @@ -125,6 +125,62 @@ struct IndexHNSWFlat : IndexHNSW { IndexHNSWFlat(int d, int M, MetricType metric = METRIC_L2); }; +/** Panorama implementation of IndexHNSWFlat following + * https://www.arxiv.org/pdf/2510.00566. + * + * Unlike cluster-based Panorama, the vectors have to be higher dimensional + * (i.e. typically d > 512) and/or be able to compress a lot of their energy in + * the early dimensions to be effective. This is because HNSW accesses vectors + * in a random order, which makes cache misses dominate the distance computation + * time. + * + * The `num_panorama_levels` parameter controls the granularity of progressive + * distance refinement, allowing candidates to be eliminated early using partial + * distance computations rather than computing full distances. + * + * NOTE: This version of HNSW handles search slightly differently than the + * vanilla HNSW, as it uses partial distance computations with progressive + * refinement bounds. Instead of computing full distances immediately for all + * candidates, Panorama maintains lower and upper bounds that are incrementally + * tightened across refinement levels. Candidates are inserted into the search + * beam using approximate distance estimates (LB+UB)/2 and are only fully + * evaluated when they survive pruning and enter the result heap. This allows + * the algorithm to prune unpromising candidates early using Cauchy-Schwarz + * bounds on partial inner products. Hence, recall is not guaranteed to be the + * same as vanilla HNSW due to the heterogeneous precision within the search + * beam (exact vs. partial distance estimates affecting traversal order). + */ +struct IndexHNSWFlatPanorama : IndexHNSWFlat { + IndexHNSWFlatPanorama(); + IndexHNSWFlatPanorama( + int d, + int M, + int num_panorama_levels, + MetricType metric = METRIC_L2); + + void add(idx_t n, const float* x) override; + void reset() override; + void permute_entries(const idx_t* perm) override; + + /// Inline for performance - called frequently in search hot path. + const float* get_cum_sum(idx_t i) const { + return cum_sums.data() + i * (num_panorama_levels + 1); + } + + /// Compute cumulative sums for a vector (used both for database points and + /// queries). + static void compute_cum_sums( + const float* x, + float* dst_cum_sums, + int d, + int num_panorama_levels, + int panorama_level_width); + + std::vector cum_sums; + const size_t panorama_level_width; + const size_t num_panorama_levels; +}; + /** PQ index topped with with a HNSW structure to access elements * more efficiently. */ diff --git a/faiss/IndexPQ.cpp b/faiss/IndexPQ.cpp index 4f7a2d0f6..255900ced 100644 --- a/faiss/IndexPQ.cpp +++ b/faiss/IndexPQ.cpp @@ -81,6 +81,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer { const float* sdc; std::vector precomputed_table; size_t ndis; + const float* q; float distance_to_code(const uint8_t* code) final { ndis++; @@ -109,7 +110,8 @@ struct PQDistanceComputer : FlatCodesDistanceComputer { : FlatCodesDistanceComputer( storage.codes.data(), storage.code_size), - pq(storage.pq) { + pq(storage.pq), + q(nullptr) { precomputed_table.resize(pq.M * pq.ksub); nb = storage.ntotal; d = storage.d; @@ -123,6 +125,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer { } void set_query(const float* x) override { + q = x; if (metric == METRIC_L2) { pq.compute_distance_table(x, precomputed_table.data()); } else { diff --git a/faiss/clone_index.cpp b/faiss/clone_index.cpp index 062165381..d62e72024 100644 --- a/faiss/clone_index.cpp +++ b/faiss/clone_index.cpp @@ -129,6 +129,7 @@ IndexIDMap* clone_IndexIDMap(const IndexIDMap* im) { IndexHNSW* clone_IndexHNSW(const IndexHNSW* ihnsw) { TRYCLONE(IndexHNSW2Level, ihnsw) + TRYCLONE(IndexHNSWFlatPanorama, ihnsw) TRYCLONE(IndexHNSWFlat, ihnsw) TRYCLONE(IndexHNSWPQ, ihnsw) TRYCLONE(IndexHNSWSQ, ihnsw) diff --git a/faiss/impl/DistanceComputer.h b/faiss/impl/DistanceComputer.h index 460e872d7..be88958b1 100644 --- a/faiss/impl/DistanceComputer.h +++ b/faiss/impl/DistanceComputer.h @@ -113,18 +113,89 @@ struct FlatCodesDistanceComputer : DistanceComputer { const uint8_t* codes; size_t code_size; - FlatCodesDistanceComputer(const uint8_t* codes, size_t code_size) - : codes(codes), code_size(code_size) {} + const float* q = nullptr; // not used in all distance computers - FlatCodesDistanceComputer() : codes(nullptr), code_size(0) {} + FlatCodesDistanceComputer( + const uint8_t* codes, + size_t code_size, + const float* q = nullptr) + : codes(codes), code_size(code_size), q(q) {} + + explicit FlatCodesDistanceComputer(const float* q) + : codes(nullptr), code_size(0), q(q) {} + + FlatCodesDistanceComputer() : codes(nullptr), code_size(0), q(nullptr) {} float operator()(idx_t i) override { return distance_to_code(codes + i * code_size); } + /// Computes a partial dot product over a slice of the query vector. + /// The slice is defined by the following parameters: + /// — `offset`: the starting index of the first component to include + /// — `num_components`: the number of consecutive components to include + /// + /// Components refer to raw dimensions of the flat (uncompressed) query + /// vector. + /// + /// By default, this method throws an error, as it is only implemented + /// in specific subclasses such as `FlatL2Dis`. Other flat distance + /// computers may override this when partial dot product support is needed. + /// + /// Over time, this method might be changed to a pure virtual function (`= + /// 0`) to enforce implementation in subclasses that require this + /// functionality. + /// + /// This method is not part of the generic `DistanceComputer` interface + /// because for compressed representations (e.g., product quantization), + /// calling `partial_dot_product` repeatedly is often less efficient than + /// computing the full distance at once. + /// + /// Supporting efficient partial scans generally requires a different memory + /// layout, such as interleaved blocks that keep SIMD lanes full. This is a + /// non-trivial change and not supported in the current flat layout. + /// + /// For more details on partial (or chunked) dot product computations and + /// the performance trade-offs involved, refer to the Panorama paper: + /// https://arxiv.org/pdf/2510.00566 + virtual float partial_dot_product( + const idx_t /* i */, + const uint32_t /* offset */, + const uint32_t /* num_components */) { + FAISS_THROW_MSG("partial_dot_product not implemented"); + } + /// compute distance of current query to an encoded vector virtual float distance_to_code(const uint8_t* code) = 0; + /// Compute partial dot products of current query to 4 stored vectors. + /// See `partial_dot_product` for more details. + virtual void partial_dot_product_batch_4( + const idx_t idx0, + const idx_t idx1, + const idx_t idx2, + const idx_t idx3, + float& dp0, + float& dp1, + float& dp2, + float& dp3, + const uint32_t offset, + const uint32_t num_components) { + // default implementation for correctness + const float d0 = + this->partial_dot_product(idx0, offset, num_components); + const float d1 = + this->partial_dot_product(idx1, offset, num_components); + const float d2 = + this->partial_dot_product(idx2, offset, num_components); + const float d3 = + this->partial_dot_product(idx3, offset, num_components); + dp0 = d0; + dp1 = d1; + dp2 = d2; + dp3 = d3; + } + virtual ~FlatCodesDistanceComputer() override {} }; diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp index e5a43f2ec..465297fe5 100644 --- a/faiss/impl/HNSW.cpp +++ b/faiss/impl/HNSW.cpp @@ -9,6 +9,8 @@ #include +#include + #include #include #include @@ -588,6 +590,28 @@ void HNSW::add_with_locks( using MinimaxHeap = HNSW::MinimaxHeap; using Node = HNSW::Node; using C = HNSW::C; + +/** Helper to extract search parameters from HNSW and SearchParameters */ +static inline void extract_search_params( + const HNSW& hnsw, + const SearchParameters* params, + bool& do_dis_check, + int& efSearch, + const IDSelector*& sel) { + // can be overridden by search params + do_dis_check = hnsw.check_relative_distance; + efSearch = hnsw.efSearch; + sel = nullptr; + if (params) { + if (const SearchParametersHNSW* hnsw_params = + dynamic_cast(params)) { + do_dis_check = hnsw_params->check_relative_distance; + efSearch = hnsw_params->efSearch; + } + sel = params->sel; + } +} + /** Do a BFS on the candidates list */ int search_from_candidates( const HNSW& hnsw, @@ -602,18 +626,10 @@ int search_from_candidates( int nres = nres_in; int ndis = 0; - // can be overridden by search params - bool do_dis_check = hnsw.check_relative_distance; - int efSearch = hnsw.efSearch; - const IDSelector* sel = nullptr; - if (params) { - if (const SearchParametersHNSW* hnsw_params = - dynamic_cast(params)) { - do_dis_check = hnsw_params->check_relative_distance; - efSearch = hnsw_params->efSearch; - } - sel = params->sel; - } + bool do_dis_check; + int efSearch; + const IDSelector* sel; + extract_search_params(hnsw, params, do_dis_check, efSearch, sel); C::T threshold = res.threshold; for (int i = 0; i < candidates.size(); i++) { @@ -735,6 +751,253 @@ int search_from_candidates( return nres; } +int search_from_candidates_panorama( + const HNSW& hnsw, + const IndexHNSW* index, + DistanceComputer& qdis, + ResultHandler& res, + MinimaxHeap& candidates, + VisitedTable& vt, + HNSWStats& stats, + int level, + int nres_in, + const SearchParameters* params) { + int nres = nres_in; + int ndis = 0; + + bool do_dis_check; + int efSearch; + const IDSelector* sel; + extract_search_params(hnsw, params, do_dis_check, efSearch, sel); + + C::T threshold = res.threshold; + for (int i = 0; i < candidates.size(); i++) { + idx_t v1 = candidates.ids[i]; + float d = candidates.dis[i]; + FAISS_ASSERT(v1 >= 0); + if (!sel || sel->is_member(v1)) { + if (d < threshold) { + if (res.add_result(d, v1)) { + threshold = res.threshold; + } + } + } + vt.set(v1); + } + + // Validate the index type so we can access cumulative sums, n_levels, and + // get the ability to compute partial dot products. + const auto* panorama_index = + dynamic_cast(index); + FAISS_THROW_IF_NOT_MSG( + panorama_index, "Index must be a IndexHNSWFlatPanorama"); + auto* flat_codes_qdis = dynamic_cast(&qdis); + FAISS_THROW_IF_NOT_MSG( + flat_codes_qdis, + "DistanceComputer must be a FlatCodesDistanceComputer"); + + // Allocate space for the index array and exact distances. + size_t M = hnsw.nb_neighbors(0); + std::vector index_array(M); + std::vector exact_distances(M); + + const float* query = flat_codes_qdis->q; + std::vector query_cum_sums(panorama_index->num_panorama_levels + 1); + IndexHNSWFlatPanorama::compute_cum_sums( + query, + query_cum_sums.data(), + panorama_index->d, + panorama_index->num_panorama_levels, + panorama_index->panorama_level_width); + float query_norm_sq = query_cum_sums[0] * query_cum_sums[0]; + + int nstep = 0; + + while (candidates.size() > 0) { + float d0 = 0; + int v0 = candidates.pop_min(&d0); + + if (do_dis_check) { + // tricky stopping condition: there are more than ef + // distances that are processed already that are smaller + // than d0 + + int n_dis_below = candidates.count_below(d0); + if (n_dis_below >= efSearch) { + break; + } + } + + size_t begin, end; + hnsw.neighbor_range(v0, level, &begin, &end); + + // Unlike the vanilla HNSW, we already remove (and compact) the visited + // nodes from the candidates list at this stage. We also remove nodes + // that are not selected. + size_t initial_size = 0; + for (size_t j = begin; j < end; j++) { + int v1 = hnsw.neighbors[j]; + if (v1 < 0) { + break; + } + + const float* cum_sums_v1 = panorama_index->get_cum_sum(v1); + index_array[initial_size] = v1; + exact_distances[initial_size] = + query_norm_sq + cum_sums_v1[0] * cum_sums_v1[0]; + + bool is_selected = !sel || sel->is_member(v1); + initial_size += is_selected && !vt.get(v1) ? 1 : 0; + + vt.set(v1); + } + + size_t batch_size = initial_size; + size_t curr_panorama_level = 0; + const size_t num_panorama_levels = panorama_index->num_panorama_levels; + while (curr_panorama_level < num_panorama_levels && batch_size > 0) { + float query_cum_norm = query_cum_sums[curr_panorama_level + 1]; + + const size_t panorama_level_width = + panorama_index->panorama_level_width; + size_t start_dim = curr_panorama_level * panorama_level_width; + size_t end_dim = (curr_panorama_level + 1) * panorama_level_width; + end_dim = std::min(end_dim, static_cast(panorama_index->d)); + + size_t i = 0; + size_t next_batch_size = 0; + for (; i + 3 < batch_size; i += 4) { + idx_t idx_0 = index_array[i]; + idx_t idx_1 = index_array[i + 1]; + idx_t idx_2 = index_array[i + 2]; + idx_t idx_3 = index_array[i + 3]; + + float dp[4]; + flat_codes_qdis->partial_dot_product_batch_4( + idx_0, + idx_1, + idx_2, + idx_3, + dp[0], + dp[1], + dp[2], + dp[3], + start_dim, + end_dim - start_dim); + ndis += 4; + + float new_exact_0 = exact_distances[i + 0] - 2 * dp[0]; + float new_exact_1 = exact_distances[i + 1] - 2 * dp[1]; + float new_exact_2 = exact_distances[i + 2] - 2 * dp[2]; + float new_exact_3 = exact_distances[i + 3] - 2 * dp[3]; + + float cum_sum_0 = panorama_index->get_cum_sum( + idx_0)[curr_panorama_level + 1]; + float cum_sum_1 = panorama_index->get_cum_sum( + idx_1)[curr_panorama_level + 1]; + float cum_sum_2 = panorama_index->get_cum_sum( + idx_2)[curr_panorama_level + 1]; + float cum_sum_3 = panorama_index->get_cum_sum( + idx_3)[curr_panorama_level + 1]; + + float cs_bound_0 = 2.0f * cum_sum_0 * query_cum_norm; + float cs_bound_1 = 2.0f * cum_sum_1 * query_cum_norm; + float cs_bound_2 = 2.0f * cum_sum_2 * query_cum_norm; + float cs_bound_3 = 2.0f * cum_sum_3 * query_cum_norm; + + float lower_bound_0 = new_exact_0 - cs_bound_0; + float lower_bound_1 = new_exact_1 - cs_bound_1; + float lower_bound_2 = new_exact_2 - cs_bound_2; + float lower_bound_3 = new_exact_3 - cs_bound_3; + + // The following code is not the most branch friendly (due to + // the maintenance of the candidate heap), but micro-benchmarks + // have shown that it is not worth it to write horrible code to + // squeeze out those cycles. + if (lower_bound_0 <= threshold) { + exact_distances[next_batch_size] = new_exact_0; + index_array[next_batch_size] = idx_0; + next_batch_size += 1; + } else { + candidates.push(idx_0, new_exact_0); + } + if (lower_bound_1 <= threshold) { + exact_distances[next_batch_size] = new_exact_1; + index_array[next_batch_size] = idx_1; + next_batch_size += 1; + } else { + candidates.push(idx_1, new_exact_1); + } + if (lower_bound_2 <= threshold) { + exact_distances[next_batch_size] = new_exact_2; + index_array[next_batch_size] = idx_2; + next_batch_size += 1; + } else { + candidates.push(idx_2, new_exact_2); + } + if (lower_bound_3 <= threshold) { + exact_distances[next_batch_size] = new_exact_3; + index_array[next_batch_size] = idx_3; + next_batch_size += 1; + } else { + candidates.push(idx_3, new_exact_3); + } + } + + // Process the remaining candidates. + for (; i < batch_size; i++) { + idx_t idx = index_array[i]; + + float dp = flat_codes_qdis->partial_dot_product( + idx, start_dim, end_dim - start_dim); + ndis += 1; + float new_exact = exact_distances[i] - 2.0f * dp; + + float cum_sum = panorama_index->get_cum_sum( + idx)[curr_panorama_level + 1]; + float cs_bound = 2.0f * cum_sum * query_cum_norm; + float lower_bound = new_exact - cs_bound; + + if (lower_bound <= threshold) { + exact_distances[next_batch_size] = new_exact; + index_array[next_batch_size] = idx; + next_batch_size += 1; + } else { + candidates.push(idx, new_exact); + } + } + + batch_size = next_batch_size; + curr_panorama_level++; + } + + // Add surviving candidates to the result handler. + for (size_t i = 0; i < batch_size; i++) { + idx_t idx = index_array[i]; + if (res.add_result(exact_distances[i], idx)) { + nres += 1; + } + candidates.push(idx, exact_distances[i]); + } + + nstep++; + if (!do_dis_check && nstep > efSearch) { + break; + } + } + + if (level == 0) { + stats.n1++; + if (candidates.size() == 0) { + stats.n2++; + } + stats.ndis += ndis; + stats.nhops += nstep; + } + + return nres; +} + std::priority_queue search_from_candidate_unbounded( const HNSW& hnsw, const Node& node, @@ -936,6 +1199,7 @@ int extract_k_from_ResultHandler(ResultHandler& res) { HNSWStats HNSW::search( DistanceComputer& qdis, + const IndexHNSW* index, ResultHandler& res, VisitedTable& vt, const SearchParameters* params) const { @@ -966,13 +1230,28 @@ HNSWStats HNSW::search( } int ef = std::max(efSearch, k); - if (bounded_queue) { // this is the most common branch + if (bounded_queue) { // this is the most common branch, for now we only + // support Panorama search in this branch MinimaxHeap candidates(ef); candidates.push(nearest, d_nearest); - search_from_candidates( - *this, qdis, res, candidates, vt, stats, 0, 0, params); + if (!is_panorama) { + search_from_candidates( + *this, qdis, res, candidates, vt, stats, 0, 0, params); + } else { + search_from_candidates_panorama( + *this, + index, + qdis, + res, + candidates, + vt, + stats, + 0, + 0, + params); + } } else { std::priority_queue top_candidates = search_from_candidate_unbounded( diff --git a/faiss/impl/HNSW.h b/faiss/impl/HNSW.h index 1a50c2a09..32ea201b1 100644 --- a/faiss/impl/HNSW.h +++ b/faiss/impl/HNSW.h @@ -8,12 +8,12 @@ #pragma once #include -#include #include #include #include +#include #include #include #include @@ -22,6 +22,10 @@ namespace faiss { +// Forward declarations to avoid circular dependency. +struct IndexHNSW; +struct IndexHNSWFlatPanorama; + /** Implementation of the Hierarchical Navigable Small World * datastructure. * @@ -146,6 +150,9 @@ struct HNSW { /// use bounded queue during exploration bool search_bounded_queue = true; + /// use Panorama progressive pruning in search + bool is_panorama = false; + // methods that initialize the tree sizes /// initialize the assign_probas and cum_nneighbor_per_level to @@ -196,9 +203,15 @@ struct HNSW { VisitedTable& vt, bool keep_max_size_level0 = false); - /// search interface for 1 point, single thread + /// Search interface for 1 point, single thread + /// + /// NOTE: We pass a reference to the index itself to allow for additional + /// state information to be passed (used for Panorama progressive pruning). + /// The alternative would be to override both HNSW::search and + /// HNSWIndex::search, which would be a nuisance of code duplication. HNSWStats search( DistanceComputer& qdis, + const IndexHNSW* index, ResultHandler& res, VisitedTable& vt, const SearchParameters* params = nullptr) const; @@ -267,6 +280,22 @@ int search_from_candidates( int nres_in = 0, const SearchParameters* params = nullptr); +/// Equivalent to `search_from_candidates`, but applies pruning with progressive +/// refinement bounds. +/// This is used in `IndexHNSWFlatPanorama` to improve the search performance +/// for higher dimensional vectors. +int search_from_candidates_panorama( + const HNSW& hnsw, + const IndexHNSW* index, + DistanceComputer& qdis, + ResultHandler& res, + HNSW::MinimaxHeap& candidates, + VisitedTable& vt, + HNSWStats& stats, + int level, + int nres_in = 0, + const SearchParameters* params = nullptr); + HNSWStats greedy_update_nearest( const HNSW& hnsw, DistanceComputer& qdis, diff --git a/faiss/impl/RaBitQuantizer.cpp b/faiss/impl/RaBitQuantizer.cpp index 79dec0593..2e1633bdd 100644 --- a/faiss/impl/RaBitQuantizer.cpp +++ b/faiss/impl/RaBitQuantizer.cpp @@ -367,6 +367,7 @@ float RaBitQDistanceComputerNotQ::distance_to_code_full(const uint8_t* code) { } void RaBitQDistanceComputerNotQ::set_query(const float* x) { + q = x; FAISS_ASSERT(x != nullptr); FAISS_ASSERT( (metric_type == MetricType::METRIC_L2 || @@ -564,6 +565,7 @@ float RaBitQDistanceComputerQ::distance_to_code_full(const uint8_t* code) { using rabitq_utils::Z_MAX_BY_QB; void RaBitQDistanceComputerQ::set_query(const float* x) { + q = x; FAISS_ASSERT(x != nullptr); FAISS_ASSERT( (metric_type == MetricType::METRIC_L2 || diff --git a/faiss/impl/ScalarQuantizer.h b/faiss/impl/ScalarQuantizer.h index 93b856dcc..1408efda3 100644 --- a/faiss/impl/ScalarQuantizer.h +++ b/faiss/impl/ScalarQuantizer.h @@ -98,9 +98,7 @@ struct ScalarQuantizer : Quantizer { SQuantizer* select_quantizer() const; struct SQDistanceComputer : FlatCodesDistanceComputer { - const float* q; - - SQDistanceComputer() : q(nullptr) {} + SQDistanceComputer() : FlatCodesDistanceComputer(nullptr) {} virtual float query_to_code(const uint8_t* code) const = 0; diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index f92ac2588..1a137e0d4 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -1168,11 +1168,15 @@ Index* read_index(IOReader* f, int io_flags) { idx = idxp; } else if ( h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") || - h == fourcc("IHN2") || h == fourcc("IHNc") || h == fourcc("IHc2")) { + h == fourcc("IHN2") || h == fourcc("IHNc") || h == fourcc("IHc2") || + h == fourcc("IHfP")) { IndexHNSW* idxhnsw = nullptr; if (h == fourcc("IHNf")) { idxhnsw = new IndexHNSWFlat(); } + if (h == fourcc("IHfP")) { + idxhnsw = new IndexHNSWFlatPanorama(); + } if (h == fourcc("IHNp")) { idxhnsw = new IndexHNSWPQ(); } @@ -1189,6 +1193,15 @@ Index* read_index(IOReader* f, int io_flags) { idxhnsw = new IndexHNSWCagra(); } read_index_header(idxhnsw, f); + if (h == fourcc("IHfP")) { + auto idx_panorama = dynamic_cast(idxhnsw); + size_t nlevels; + READ1(nlevels); + const_cast(idx_panorama->num_panorama_levels) = nlevels; + const_cast(idx_panorama->panorama_level_width) = + (idx_panorama->d + nlevels - 1) / nlevels; + READVECTOR(idx_panorama->cum_sums); + } if (h == fourcc("IHNc") || h == fourcc("IHc2")) { READ1(idxhnsw->keep_max_size_level0); auto idx_hnsw_cagra = dynamic_cast(idxhnsw); @@ -1201,6 +1214,7 @@ Index* read_index(IOReader* f, int io_flags) { } } read_HNSW(&idxhnsw->hnsw, f); + idxhnsw->hnsw.is_panorama = (h == fourcc("IHfP")); idxhnsw->storage = read_index(f, io_flags); idxhnsw->own_fields = idxhnsw->storage != nullptr; if (h == fourcc("IHNp") && !(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { @@ -1603,6 +1617,7 @@ IndexBinary* read_index_binary(IOReader* f, int io_flags) { IndexBinaryHNSW* idxhnsw = new IndexBinaryHNSW(); read_index_binary_header(idxhnsw, f); read_HNSW(&idxhnsw->hnsw, f); + idxhnsw->hnsw.is_panorama = false; idxhnsw->storage = read_index_binary(f, io_flags); idxhnsw->own_fields = true; idx = idxhnsw; @@ -1613,6 +1628,7 @@ IndexBinary* read_index_binary(IOReader* f, int io_flags) { READ1(idxhnsw->base_level_only); READ1(idxhnsw->num_base_level_search_entrypoints); read_HNSW(&idxhnsw->hnsw, f); + idxhnsw->hnsw.is_panorama = false; idxhnsw->storage = read_index_binary(f, io_flags); idxhnsw->own_fields = true; idx = idxhnsw; diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index fec6d8168..f98dadebc 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -832,15 +832,23 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index(idxmap->index, f); WRITEVECTOR(idxmap->id_map); } else if (const IndexHNSW* idxhnsw = dynamic_cast(idx)) { - uint32_t h = dynamic_cast(idx) ? fourcc("IHNf") - : dynamic_cast(idx) ? fourcc("IHNp") - : dynamic_cast(idx) ? fourcc("IHNs") - : dynamic_cast(idx) ? fourcc("IHN2") - : dynamic_cast(idx) ? fourcc("IHc2") - : 0; + uint32_t h = dynamic_cast(idx) + ? fourcc("IHfP") + : dynamic_cast(idx) ? fourcc("IHNf") + : dynamic_cast(idx) ? fourcc("IHNp") + : dynamic_cast(idx) ? fourcc("IHNs") + : dynamic_cast(idx) ? fourcc("IHN2") + : dynamic_cast(idx) ? fourcc("IHc2") + : 0; FAISS_THROW_IF_NOT(h != 0); WRITE1(h); write_index_header(idxhnsw, f); + if (h == fourcc("IHfP")) { + auto idx_panorama = + dynamic_cast(idxhnsw); + WRITE1(idx_panorama->num_panorama_levels); + WRITEVECTOR(idx_panorama->cum_sums); + } if (h == fourcc("IHc2")) { WRITE1(idxhnsw->keep_max_size_level0); auto idx_hnsw_cagra = dynamic_cast(idxhnsw); diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index add0c2927..547b7053e 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -493,6 +493,11 @@ IndexHNSW* parse_IndexHNSW( return new IndexHNSWFlat(d, hnsw_M, mt); } + if (match("FlatPanorama([0-9]+)?")) { + int nlevels = mres_to_int(sm[1], 8); // default to 8 levels + return new IndexHNSWFlatPanorama(d, hnsw_M, nlevels, mt); + } + if (match("PQ([0-9]+)(x[0-9]+)?(np)?")) { int M = std::stoi(sm[1].str()); int nbit = mres_to_int(sm[2], 8, 1); diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index de48dff1e..c034ef390 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -782,6 +782,7 @@ struct faiss::simd16uint16 {}; DOWNCAST ( IndexLattice ) DOWNCAST ( IndexPreTransform ) DOWNCAST ( MultiIndexQuantizer ) + DOWNCAST ( IndexHNSWFlatPanorama ) DOWNCAST ( IndexHNSWFlat ) DOWNCAST ( IndexHNSWPQ ) DOWNCAST ( IndexHNSWSQ ) diff --git a/tests/test_hnsw_panorama.py b/tests/test_hnsw_panorama.py new file mode 100644 index 000000000..ada7a560a --- /dev/null +++ b/tests/test_hnsw_panorama.py @@ -0,0 +1,523 @@ +# 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. + +""" +Comprehensive test suite for IndexHNSWFlatPanorama. + +IndexHNSWFlatPanorama combines the hierarchical navigable small world (HNSW) +graph structure with Panorama's level-oriented storage and progressive +filtering for efficient approximate nearest neighbor search. + +Paper: https://www.arxiv.org/pdf/2510.00566 +""" + +import unittest +import tempfile +import os +import numpy as np + +import faiss +from faiss.contrib.datasets import SyntheticDataset + + +class TestIndexHNSWFlatPanorama(unittest.TestCase): + """Test Suite for IndexHNSWFlatPanorama.""" + + # Helper methods for index creation and data generation + + def generate_data(self, d, nt, nb, nq, seed=1234): + """Generate synthetic data using SyntheticDataset.""" + ds = SyntheticDataset(d, nt, nb, nq, seed=seed) + return ds.get_train(), ds.get_database(), ds.get_queries() + + def compute_ground_truth(self, xb, xq, k): + """Compute ground truth using brute force search.""" + d = xb.shape[1] + index = faiss.IndexFlatL2(d) + index.add(xb) + D, I = index.search(xq, k) + return D, I + + def compute_recall(self, gt_I, test_I): + """Compute recall@k - fraction of ground truth results found in test results.""" + nq, k = gt_I.shape + recalls = [ + np.isin(gt_I[i], test_I[i]).sum() + for i in range(nq) + ] + return sum(recalls) / (nq * k) + + # Core functionality tests + + def test_basic_construction(self): + """Test basic index construction and properties.""" + d = 128 + M = 16 + num_levels = 8 + + index = faiss.IndexHNSWFlatPanorama(d, M, num_levels) + + self.assertEqual(index.d, d) + self.assertEqual(index.ntotal, 0) + self.assertTrue(index.is_trained) + self.assertEqual(index.hnsw.nb_neighbors(0), 2 * M) # Level 0 has 2*M neighbors + self.assertEqual(index.num_panorama_levels, num_levels) + self.assertEqual(index.metric_type, faiss.METRIC_L2) + + def test_add_and_search(self): + """Test basic add and search operations.""" + d = 128 + nb = 1000 + nt = 1500 + nq = 10 + k = 10 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=1234) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + index.add(xb) + + self.assertEqual(index.ntotal, nb) + self.assertEqual(index.cum_sums.size(), nb * (8 + 1)) + + # Search + D, I = index.search(xq, k) + + # Check that results are valid + self.assertTrue(np.all(I >= -1)) + self.assertTrue(np.all(I < nb)) + self.assertTrue(np.all(D >= 0.0)) + + def test_recall(self): + """Test search recall quality.""" + d = 128 + nb = 1000 + nt = 1500 + nq = 100 + k = 10 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=42) + + # Create index + index = faiss.IndexHNSWFlatPanorama(d, 32, 8) + index.hnsw.efConstruction = 40 + index.hnsw.efSearch = 64 + + # Add to index + index.add(xb) + + # Search + D, I = index.search(xq, k) + + # Compute ground truth + gt_D, gt_I = self.compute_ground_truth(xb, xq, k) + + # Check recall + recall = self.compute_recall(gt_I, I) + print(f"Recall@{k}: {recall}") + + # With efSearch=64, we should get reasonably good recall + # The threshold is lower than vanilla HNSW because of approximate distances + self.assertGreaterEqual(recall, 0.85) + + def test_different_panorama_levels(self): + """Test with different numbers of panorama levels.""" + d = 256 + nb = 500 + nt = 700 + nq = 50 + k = 5 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=456) + + # Compute ground truth + gt_D, gt_I = self.compute_ground_truth(xb, xq, k) + + # Test different number of panorama levels + for nlevels in [4, 8, 16]: + with self.subTest(nlevels=nlevels): + index = faiss.IndexHNSWFlatPanorama(d, 32, nlevels) + index.hnsw.efSearch = 64 + + index.add(xb) + + D, I = index.search(xq, k) + + recall = self.compute_recall(gt_I, I) + print(f"Recall@{k} with {nlevels} levels: {recall}") + + # More levels should still maintain reasonable recall + self.assertGreaterEqual(recall, 0.80) + + def test_consistency(self): + """Test that search results are consistent across multiple searches.""" + d = 64 + nb = 500 + nt = 700 + nq = 10 + k = 5 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=123) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + index.add(xb) + + # Search twice and compare results + D1, I1 = index.search(xq, k) + D2, I2 = index.search(xq, k) + + # Results should be identical + np.testing.assert_array_equal(I1, I2) + np.testing.assert_array_equal(D1, D2) + + def test_io(self): + """Test serialization and deserialization.""" + d = 64 + nb = 500 + nt = 700 + nq = 10 + k = 5 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=2024) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + index.add(xb) + + # Get search results before saving + D_before, I_before = index.search(xq, k) + + # Save and load using temporary file + fname = tempfile.mktemp(suffix='.index') + try: + faiss.write_index(index, fname) + loaded_index = faiss.read_index(fname) + + self.assertIsInstance(loaded_index, faiss.IndexHNSWFlatPanorama) + self.assertEqual(loaded_index.d, d) + self.assertEqual(loaded_index.ntotal, nb) + self.assertEqual(loaded_index.num_panorama_levels, 8) + + # Search after loading + D_after, I_after = loaded_index.search(xq, k) + + # Results should be identical after serialization + np.testing.assert_array_equal(I_before, I_after) + np.testing.assert_array_equal(D_before, D_after) + finally: + if os.path.exists(fname): + os.unlink(fname) + + def test_clone(self): + """Test index cloning.""" + d = 64 + nb = 500 + nt = 700 + nq = 10 + k = 5 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=2025) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + index.add(xb) + + # Get search results before cloning + D_before, I_before = index.search(xq, k) + + # Clone + cloned_index = faiss.clone_index(index) + + self.assertIsInstance(cloned_index, faiss.IndexHNSWFlatPanorama) + self.assertEqual(cloned_index.d, d) + self.assertEqual(cloned_index.ntotal, nb) + self.assertEqual(cloned_index.num_panorama_levels, 8) + + # Search after cloning + D_after, I_after = cloned_index.search(xq, k) + + # Results should be identical after cloning + np.testing.assert_array_equal(I_before, I_after) + np.testing.assert_array_equal(D_before, D_after) + + def test_factory(self): + """Test factory string creation.""" + d = 64 + nb = 500 + nt = 700 + + # Test factory creation with default levels + index1 = faiss.index_factory(d, "HNSW32_FlatPanorama") + self.assertIsInstance(index1, faiss.IndexHNSWFlatPanorama) + self.assertEqual(index1.d, d) + self.assertEqual(index1.hnsw.nb_neighbors(0), 2 * 32) # Level 0 has 2*M neighbors + self.assertEqual(index1.num_panorama_levels, 8) # default + + # Test factory creation with explicit levels + index2 = faiss.index_factory(d, "HNSW16_FlatPanorama12") + self.assertIsInstance(index2, faiss.IndexHNSWFlatPanorama) + self.assertEqual(index2.d, d) + self.assertEqual(index2.hnsw.nb_neighbors(0), 2 * 16) # Level 0 has 2*M neighbors + self.assertEqual(index2.num_panorama_levels, 12) + + # Test that it works + _, xb, _ = self.generate_data(d, nt, nb, 10, seed=888) + index2.add(xb) + self.assertEqual(index2.ntotal, nb) + + def test_reset(self): + """Test index reset.""" + d = 64 + nb = 500 + nt = 700 + + # Generate data + _, xb, _ = self.generate_data(d, nt, nb, 10, seed=999) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + index.add(xb) + + self.assertEqual(index.ntotal, nb) + self.assertGreater(index.cum_sums.size(), 0) + + # Reset + index.reset() + + self.assertEqual(index.ntotal, 0) + self.assertEqual(index.cum_sums.size(), 0) + + def test_high_dimensional(self): + """Test with high-dimensional data (Panorama is designed for this).""" + d = 512 + nb = 500 + nt = 700 + nq = 50 + k = 10 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=789) + + index = faiss.IndexHNSWFlatPanorama(d, 32, 16) + index.hnsw.efSearch = 128 + index.add(xb) + + D, I = index.search(xq, k) + + # Compute ground truth + gt_D, gt_I = self.compute_ground_truth(xb, xq, k) + + recall = self.compute_recall(gt_I, I) + print(f"High-dimensional (d={d}) Recall@{k}: {recall}") + + # Should work well for high-dimensional data + self.assertGreaterEqual(recall, 0.85) + + def test_add_after_search(self): + """Test adding vectors after initial search.""" + d = 64 + nb_initial = 300 + nt_initial = 500 + nb_add = 200 + nt_add = 300 + nq = 10 + k = 5 + + # Generate data + _, xb_initial, xq = self.generate_data(d, nt_initial, nb_initial, nq, seed=1111) + _, xb_add, _ = self.generate_data(d, nt_add, nb_add, 1, seed=2222) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + + # Add initial vectors + index.add(xb_initial) + + # Search on initial index + D_before, I_before = index.search(xq, k) + + # Add more vectors after search + index.add(xb_add) + + self.assertEqual(index.ntotal, nb_initial + nb_add) + self.assertEqual( + index.cum_sums.size(), + (nb_initial + nb_add) * (index.num_panorama_levels + 1) + ) + + # Search again - should work correctly with all vectors + D_after, I_after = index.search(xq, k) + + # Compute ground truth on all vectors + xb_all = np.vstack([xb_initial, xb_add]) + + gt_D, gt_I = self.compute_ground_truth(xb_all, xq, k) + + recall = self.compute_recall(gt_I, I_after) + print(f"Recall after adding more vectors: {recall}") + + # Recall might be slightly lower than single-batch due to HNSW graph structure + self.assertGreaterEqual(recall, 0.80) + + # Verify that previously found neighbors can still be found + # (they might not be in exact same positions due to new neighbors) + found_count = 0 + for i in range(nq): + for j in range(k): + if I_before[i, j] >= 0: + # Check if this neighbor is still in the result set + if I_before[i, j] in I_after[i, :]: + found_count += 1 + + retention = float(found_count) / (nq * k) + print(f"Retention of previous neighbors: {retention}") + # Should retain a reasonable number of previous neighbors (new ones might push some out) + # The threshold is lower to account for the approximate nature of HNSW + self.assertGreaterEqual(retention, 0.5) + + def test_permute_entries(self): + """Test permuting entries in the index.""" + d = 64 + nb = 500 + nt = 700 + nq = 10 + k = 5 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=3333) + + index = faiss.IndexHNSWFlatPanorama(d, 16, 8) + index.add(xb) + + # Search before permutation + D_before, I_before = index.search(xq, k) + + # Create a permutation (reverse order for simplicity) + perm = np.arange(nb - 1, -1, -1, dtype=np.int64) + + # Store cum_sums before permutation for verification + cum_sums_before = faiss.vector_to_array(index.cum_sums).copy() + + # Apply permutation + index.permute_entries(perm) + + # Verify cum_sums were permuted correctly + cum_sums_after = faiss.vector_to_array(index.cum_sums) + for i in range(nb): + src = perm[i] + for j in range(index.num_panorama_levels + 1): + expected = cum_sums_before[src * (index.num_panorama_levels + 1) + j] + actual = cum_sums_after[i * (index.num_panorama_levels + 1) + j] + self.assertEqual( + actual, expected, + f"cum_sums not permuted correctly at i={i}, j={j}" + ) + + # Search after permutation + D_after, I_after = index.search(xq, k) + + # Results should be identical (with permuted IDs) + np.testing.assert_allclose(D_before, D_after, + err_msg="Distance changed after permutation") + + # The ID should be the permuted version + # If before we found vector j, after permutation we should find + # the new position of vector j + for i in range(nq): + for j in range(k): + if I_before[i, j] >= 0: + # Find where I_before[i, j] moved to + old_id = I_before[i, j] + new_id = np.where(perm == old_id)[0][0] + self.assertEqual( + I_after[i, j], new_id, + f"Permuted ID mismatch at position ({i}, {j})" + ) + + # Verify overall recall is maintained + gt_D, gt_I = self.compute_ground_truth(xb, xq, k) + + # Map ground truth IDs through the permutation + gt_I_permuted = np.copy(gt_I) + for i in range(nq): + for j in range(k): + if gt_I[i, j] >= 0: + # Find where gt_I[i, j] moved to + gt_I_permuted[i, j] = np.where(perm == gt_I[i, j])[0][0] + + recall = self.compute_recall(gt_I_permuted, I_after) + print(f"Recall after permutation: {recall}") + self.assertGreaterEqual(recall, 0.85) + + def test_id_selector_range(self): + """Test ID filtering with range selector.""" + d = 128 + nb = 1000 + nt = 1500 + nq = 10 + k = 10 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=321) + + index = faiss.IndexHNSWFlatPanorama(d, 32, 8) + index.hnsw.efSearch = 64 + index.add(xb) + + selector = faiss.IDSelectorRange(200, 600) + params = faiss.SearchParametersHNSW() + params.sel = selector + + D, I = index.search(xq, k, params=params) + + # Verify all returned IDs are in range (HNSW may return < k) + count = 0 + for i in range(nq): + for j in range(k): + if I[i, j] != -1: + self.assertGreaterEqual(I[i, j], 200) + self.assertLess(I[i, j], 600) + count += 1 + + print(f"IDSelectorRange test: {count} results") + + def test_id_selector_batch(self): + """Test ID filtering with batch selector.""" + d = 128 + nb = 1000 + nt = 1500 + nq = 10 + k = 5 + + # Generate data + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=654) + + index = faiss.IndexHNSWFlatPanorama(d, 32, 8) + index.hnsw.efSearch = 64 + index.add(xb) + + allowed_ids = np.arange(0, nb, 10, dtype=np.int64) + + selector = faiss.IDSelectorBatch(allowed_ids) + params = faiss.SearchParametersHNSW() + params.sel = selector + + D, I = index.search(xq, k, params=params) + + # Verify all returned IDs are in allowed set (HNSW may return < k) + allowed_set = set(allowed_ids) + count = 0 + for i in range(nq): + for j in range(k): + if I[i, j] != -1: + self.assertIn( + I[i, j], allowed_set, + f"ID {I[i, j]} not in allowed set" + ) + count += 1 + + print(f"IDSelectorBatch test: {count} results")