Measure recall against brute-force truth, and see the filtered-search collapse

# /// script
# requires-python = ">=3.10"
# dependencies = ["hnswlib", "numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure recall against brute-force truth.

An HNSW index (hnswlib) over 200,000 random Gaussian 128-dimensional vectors;
exact top 10 by brute force for 1,000 queries; recall@10 and queries per
second for efSearch 10 to 200. Then a filter that keeps 0.5% of vectors:
post-filtering the top 10, post-filtering 1,000 candidates, and pre-filtering
by brute force over the matching vectors.
"""
import argparse, json, time
import hnswlib, numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=200_000, help="Vectors.")
p.add_argument("--dim", type=int, default=128, help="Dimension.")
p.add_argument("--queries", type=int, default=1000, help="Queries.")
p.add_argument("--selectivity", type=float, default=0.005, help="Share of vectors that pass the filter.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0)
X = rng.standard_normal((args.n, args.dim), dtype=np.float32); Q = rng.standard_normal((args.queries, args.dim), dtype=np.float32)
def exact(Qs, Xs, k=10):
    d = (Xs ** 2).sum(1)[None, :] - 2 * Qs @ Xs.T
    return np.argpartition(d, k, axis=1)[:, :k]
truth = np.vstack([exact(Q[i:i + 100], X) for i in range(0, args.queries, 100)])
idx = hnswlib.Index(space="l2", dim=args.dim); idx.init_index(max_elements=args.n, ef_construction=200, M=16)
t = time.perf_counter(); idx.add_items(X, num_threads=4); build_s = time.perf_counter() - t
idx.set_num_threads(1)
def recall(found, true): return np.mean([len(set(f) & set(tr)) / len(tr) for f, tr in zip(found, true)])
curve = []
for ef in (10, 20, 50, 100, 200):
    idx.set_ef(max(ef, 10)); t = time.perf_counter(); labels, _ = idx.knn_query(Q, k=10); el = time.perf_counter() - t
    curve.append({"ef_search": ef, "recall": float(recall(labels, truth)), "qps": args.queries / el})
allowed = rng.random(args.n) < args.selectivity; ids = np.flatnonzero(allowed)
ftruth = np.vstack([ids[exact(Q[i:i + 100], X[ids])] for i in range(0, args.queries, 100)])
filt = []
idx.set_ef(200); t = time.perf_counter(); lab10, _ = idx.knn_query(Q, k=10); t10 = time.perf_counter() - t
kept10 = [[x for x in row if allowed[x]] for row in lab10]
filt.append({"method": "post-filter the top 10", "mean_results": float(np.mean([len(r) for r in kept10])), "recall": float(recall(kept10, ftruth)), "ms_per_query": t10 / args.queries * 1000})
idx.set_ef(1000); t = time.perf_counter(); lab1k, _ = idx.knn_query(Q, k=1000); t1k = time.perf_counter() - t
kept1k = [[x for x in row if allowed[x]][:10] for row in lab1k]
filt.append({"method": "post-filter 1,000 candidates", "mean_results": float(np.mean([len(r) for r in kept1k])), "recall": float(recall(kept1k, ftruth)), "ms_per_query": t1k / args.queries * 1000})
t = time.perf_counter(); pre = np.vstack([ids[exact(Q[i:i + 100], X[ids])] for i in range(0, args.queries, 100)]); tp = time.perf_counter() - t
filt.append({"method": f"pre-filter, brute force over {len(ids):,} vectors", "mean_results": 10.0, "recall": float(recall(pre, ftruth)), "ms_per_query": tp / args.queries * 1000})
summary = {"build_s": build_s, "allowed": int(len(ids))}
for r in curve: print(f"efSearch {r['ef_search']:>3}: recall@10 {r['recall']:.3f}, {r['qps']:8,.0f} queries/s")
for r in filt: print(f"{r['method']:<44} results {r['mean_results']:5.2f} of 10, recall {r['recall']:.3f}, {r['ms_per_query']:.2f} ms per query")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"curve": curve, "filtered": filt}}, f)