Choose nlist for IVF

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "faiss-cpu"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Choose nlist for an IVF vector index.

An IVF index (FAISS) over random 128-dimensional vectors for nlist from 100
to 16,000 and nprobe 1, 8 and 32: queries per second and recall@10 against
exact search on 1,000 queries.
"""
import argparse, json, time
import numpy as np, faiss

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=1_000_000, help="Vectors.")
p.add_argument("--dim", type=int, default=128, help="Dimensions.")
p.add_argument("--nlists", type=lambda s: [int(v) for v in s.split(",")], default=[100, 1000, 4000, 16000], help="nlist values.")
p.add_argument("--nprobes", type=lambda s: [int(v) for v in s.split(",")], default=[1, 8, 32], help="nprobe values.")
p.add_argument("--queries", type=int, default=1000, help="Queries.")
p.add_argument("--k", type=int, default=10, help="Neighbours.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)
xb = rng.random((args.n, args.dim), dtype=np.float32); xq = rng.random((args.queries, args.dim), dtype=np.float32)
flat = faiss.IndexFlatL2(args.dim); flat.add(xb)
_, truth = flat.search(xq, args.k)
rows = []
for nl in args.nlists:
    q = faiss.IndexFlatL2(args.dim); idx = faiss.IndexIVFFlat(q, args.dim, nl)
    idx.train(xb[: max(nl * 40, 100_000)]); idx.add(xb)
    for npb in args.nprobes:
        idx.nprobe = npb
        t = time.perf_counter(); _, got = idx.search(xq, args.k); dt = time.perf_counter() - t
        recall = float(np.mean([len(set(g) & set(tr)) / args.k for g, tr in zip(got, truth)]))
        rows.append({"nlist": nl, "nprobe": npb, "qps": args.queries / dt, "recall": recall})
        print(f"nlist {nl:>6} nprobe {npb:>3}: {args.queries / dt:9,.0f} queries/s   recall@{args.k} {recall:.3f}")
print(f"sqrt(N) = {args.n ** 0.5:.0f}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"settings": rows, "summary": [{"sqrt_n": args.n ** 0.5}]}}, f)