mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Reapply IndexHNSWFlatPanorama with backward compatible serialization (#4692)
Summary: As title, this PR reverts the prior [revert PR](https://github.com/facebookresearch/faiss/pull/4690). mnorris11 is this fix OK? Please feel free to push your test in here. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4692 Reviewed By: subhadeepkaran Differential Revision: D88390767 Pulled By: mnorris11 fbshipit-source-id: 1f57c46cfbfad5d46f2e20a2690c18bb298113d1
This commit is contained in:
Alexis Schlomer
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
337dfe8043
commit
7372bc7982
@@ -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)
|
||||
@@ -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) {
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
+58
-4
@@ -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<const float*>(codes + idx0 * code_size);
|
||||
const float* __restrict y1 =
|
||||
reinterpret_cast<const float*>(codes + idx1 * code_size);
|
||||
const float* __restrict y2 =
|
||||
reinterpret_cast<const float*>(codes + idx2 * code_size);
|
||||
const float* __restrict y3 =
|
||||
reinterpret_cast<const float*>(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 {
|
||||
|
||||
+90
-1
@@ -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<float> 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
|
||||
**************************************************************/
|
||||
|
||||
+57
-1
@@ -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<float> 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.
|
||||
*/
|
||||
|
||||
+4
-1
@@ -81,6 +81,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
|
||||
const float* sdc;
|
||||
std::vector<float> 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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {}
|
||||
};
|
||||
|
||||
|
||||
+294
-15
@@ -9,6 +9,8 @@
|
||||
|
||||
#include <cstddef>
|
||||
|
||||
#include <faiss/IndexHNSW.h>
|
||||
|
||||
#include <faiss/impl/AuxIndexStructures.h>
|
||||
#include <faiss/impl/DistanceComputer.h>
|
||||
#include <faiss/impl/IDSelector.h>
|
||||
@@ -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<const SearchParametersHNSW*>(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<const SearchParametersHNSW*>(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<C>& 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<const IndexHNSWFlatPanorama*>(index);
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
panorama_index, "Index must be a IndexHNSWFlatPanorama");
|
||||
auto* flat_codes_qdis = dynamic_cast<FlatCodesDistanceComputer*>(&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<idx_t> index_array(M);
|
||||
std::vector<float> exact_distances(M);
|
||||
|
||||
const float* query = flat_codes_qdis->q;
|
||||
std::vector<float> 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<size_t>(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<HNSW::Node> search_from_candidate_unbounded(
|
||||
const HNSW& hnsw,
|
||||
const Node& node,
|
||||
@@ -936,6 +1199,7 @@ int extract_k_from_ResultHandler(ResultHandler<C>& res) {
|
||||
|
||||
HNSWStats HNSW::search(
|
||||
DistanceComputer& qdis,
|
||||
const IndexHNSW* index,
|
||||
ResultHandler<C>& 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<Node> top_candidates =
|
||||
search_from_candidate_unbounded(
|
||||
|
||||
+31
-2
@@ -8,12 +8,12 @@
|
||||
#pragma once
|
||||
|
||||
#include <queue>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include <omp.h>
|
||||
|
||||
#include <faiss/Index.h>
|
||||
#include <faiss/impl/DistanceComputer.h>
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/maybe_owned_vector.h>
|
||||
#include <faiss/impl/platform_macros.h>
|
||||
@@ -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<C>& 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<HNSW::C>& 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,
|
||||
|
||||
@@ -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 ||
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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<IndexHNSWFlatPanorama*>(idxhnsw);
|
||||
size_t nlevels;
|
||||
READ1(nlevels);
|
||||
const_cast<size_t&>(idx_panorama->num_panorama_levels) = nlevels;
|
||||
const_cast<size_t&>(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<IndexHNSWCagra*>(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;
|
||||
|
||||
@@ -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<const IndexHNSW*>(idx)) {
|
||||
uint32_t h = dynamic_cast<const IndexHNSWFlat*>(idx) ? fourcc("IHNf")
|
||||
: dynamic_cast<const IndexHNSWPQ*>(idx) ? fourcc("IHNp")
|
||||
: dynamic_cast<const IndexHNSWSQ*>(idx) ? fourcc("IHNs")
|
||||
: dynamic_cast<const IndexHNSW2Level*>(idx) ? fourcc("IHN2")
|
||||
: dynamic_cast<const IndexHNSWCagra*>(idx) ? fourcc("IHc2")
|
||||
: 0;
|
||||
uint32_t h = dynamic_cast<const IndexHNSWFlatPanorama*>(idx)
|
||||
? fourcc("IHfP")
|
||||
: dynamic_cast<const IndexHNSWFlat*>(idx) ? fourcc("IHNf")
|
||||
: dynamic_cast<const IndexHNSWPQ*>(idx) ? fourcc("IHNp")
|
||||
: dynamic_cast<const IndexHNSWSQ*>(idx) ? fourcc("IHNs")
|
||||
: dynamic_cast<const IndexHNSW2Level*>(idx) ? fourcc("IHN2")
|
||||
: dynamic_cast<const IndexHNSWCagra*>(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<const IndexHNSWFlatPanorama*>(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<const IndexHNSWCagra*>(idxhnsw);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -782,6 +782,7 @@ struct faiss::simd16uint16 {};
|
||||
DOWNCAST ( IndexLattice )
|
||||
DOWNCAST ( IndexPreTransform )
|
||||
DOWNCAST ( MultiIndexQuantizer )
|
||||
DOWNCAST ( IndexHNSWFlatPanorama )
|
||||
DOWNCAST ( IndexHNSWFlat )
|
||||
DOWNCAST ( IndexHNSWPQ )
|
||||
DOWNCAST ( IndexHNSWSQ )
|
||||
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user