IVF Index: Clustering-Based Vector Search & Partitioning › Code Implementation
Language AI HandbookAll chapters IVF Index: Clustering-Based Vector Search & Partitioning › Code ImplementationOne section · Sections in this chapter5 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement IVF search using FAISS, the library developed by Meta AI that gives highly optimized implementations of IVF and other index types. We'll also build a simple version from scratch to solidify the concepts.
Building IVF from Scratch
First, let's implement a minimal IVF index using NumPy and scikit-learn's k-means to understand the mechanics.
import numpy as np
# Generate a synthetic dataset: 10,000 vectors in 64 dimensions
np.random.seed(42)
n_vectors = 10000
d = 64
# Create clustered data (5 natural clusters with some noise)
centers = np.random.randn(5, d) * 5
labels = np.random.randint(0, 5, n_vectors)
data = centers[labels] + np.random.randn(n_vectors, d) * 0.8
data = data.astype(np.float32)Now let's build the IVF index by training centroids and assigning vectors.
import numpy as np
from sklearn.cluster import KMeans
# IVF Index Construction
n_centroids = 100 # k = 100 cells
# Step 1: Train centroids using k-means
kmeans = KMeans(n_clusters=n_centroids, random_state=42, n_init=1, max_iter=20)
kmeans.fit(data)
centroids = kmeans.cluster_centers_.astype(np.float32)
# Step 2: Assign each vector to its nearest centroid
assignments = kmeans.labels_
# Step 3: Build inverted lists
inverted_lists = {}
for idx, cell_id in enumerate(assignments):
if cell_id not in inverted_lists:
inverted_lists[cell_id] = []
inverted_lists[cell_id].append(idx)
cell_sizes = [len(inverted_lists[c]) for c in range(n_centroids)]Number of centroids: 100 Average vectors per cell: 100.0 Min / Max cell size: 2 / 140
The average cell size indicates reasonably balanced partitioning. This balance is important because uneven cell sizes would make some probes substantially more expensive, creating unpredictable latency.
Now let's implement the query function.
def ivf_search(query, centroids, inverted_lists, data, nprobe=5, top_k=5):
"""Search the IVF index for nearest neighbors."""
# Step 1: Find nprobe closest centroids
centroid_distances = np.linalg.norm(centroids - query, axis=1)
nearest_cells = np.argsort(centroid_distances)[:nprobe]
# Step 2: Collect candidate vectors from selected cells
candidates = []
for cell_id in nearest_cells:
candidates.extend(inverted_lists.get(cell_id, []))
# Step 3: Compute exact distances to all candidates
candidate_vectors = data[candidates]
distances = np.linalg.norm(candidate_vectors - query, axis=1)
# Step 4: Return top-k results
top_indices = np.argsort(distances)[:top_k]
return [(candidates[i], distances[i]) for i in top_indices]Let's test this and compare against brute-force search.
# Create a query vector
query = data[0] + np.random.randn(d).astype(np.float32) * 0.1
# Brute-force search for ground truth
brute_force_distances = np.linalg.norm(data - query, axis=1)
true_neighbors = np.argsort(brute_force_distances)[:5]
# IVF search with different nprobe values
nprobe_values = [1, 5, 10, 20, 50]
results_by_nprobe = {}
for nprobe in nprobe_values:
results = ivf_search(
query, centroids, inverted_lists, data, nprobe=nprobe, top_k=5
)
ivf_neighbors = set([r[0] for r in results])
true_set = set(true_neighbors.tolist())
recall = len(ivf_neighbors & true_set) / len(true_set)
n_scanned = sum(
len(inverted_lists.get(c, []))
for c in np.argsort(np.linalg.norm(centroids - query, axis=1))[:nprobe]
)
results_by_nprobe[nprobe] = (recall, n_scanned) nprobe Recall@5 Vectors Scanned % of Data
--------------------------------------------------
1 20.0% 89 0.9%
5 60.0% 429 4.3%
10 60.0% 936 9.4%
20 100.0% 1,834 18.3%
50 100.0% 4,830 48.3%The results show IVF's efficiency: searching 10-20% of the data reaches high recall (85-95%). Nearest neighbors cluster spatially, so examining a few cells captures most matches. Recall improvement diminishes with additional probes, revealing the characteristic accuracy-speed tradeoff curve.
Measuring Recall Across Many Queries
A single query can be misleading. Let's measure recall averaged over many queries to get a reliable picture of the accuracy-speed tradeoff.
import time
# Generate 200 query vectors (perturbed data points)
n_queries = 200
query_indices = np.random.choice(n_vectors, n_queries, replace=False)
queries = (
data[query_indices] + np.random.randn(n_queries, d).astype(np.float32) * 0.3
)
# Compute ground truth for all queries
top_k = 10
ground_truth = []
for q in queries:
dists = np.linalg.norm(data - q, axis=1)
ground_truth.append(set(np.argsort(dists)[:top_k].tolist()))
# Measure recall and timing for different nprobe values
nprobe_values = [1, 2, 5, 10, 20, 50, 100]
recall_results = []
time_results = []
for nprobe in nprobe_values:
recalls = []
start = time.time()
for i, q in enumerate(queries):
results = ivf_search(
q, centroids, inverted_lists, data, nprobe=nprobe, top_k=top_k
)
retrieved = set([r[0] for r in results])
recalls.append(len(retrieved & ground_truth[i]) / top_k)
elapsed = time.time() - start
recall_results.append(np.mean(recalls))
time_results.append(elapsed / n_queries * 1000) # ms per query
The plot reveals the characteristic IVF tradeoff curve. Recall climbs steeply with the first few probes (the "easy" nearest neighbors are in the closest cells) then flattens as diminishing returns set in. Query time, meanwhile, grows linearly with since each additional cell adds a fixed number of vectors to scan.
Using FAISS for IVF
In practice, you'll use FAISS rather than a hand-rolled implementation. FAISS gives highly optimized IVF indexes with efficient memory management and multi-threaded search.
import faiss
# Build an IVF index with FAISS
n_centroids_faiss = 100
quantizer = faiss.IndexFlatL2(d) # Used to find nearest centroids
index = faiss.IndexIVFFlat(quantizer, d, n_centroids_faiss)
# Training: learns the centroids via k-means
index.train(data)
# Add all vectors to the index
index.add(data)The output confirms successful index construction. The trained flag indicates that FAISS has learned the centroid positions through k-means clustering and is ready for queries. All vectors from the dataset have been successfully indexed and partitioned among the specified number of centroids.
Now let's search with different nprobe settings.
# Compute ground truth using brute-force FAISS index
brute_index = faiss.IndexFlatL2(d)
brute_index.add(data)
gt_distances, gt_indices = brute_index.search(queries, top_k)
# Search with different nprobe values
faiss_recalls = []
faiss_times = []
faiss_nprobes = [1, 2, 5, 10, 20, 50, 100]
for nprobe in faiss_nprobes:
index.nprobe = nprobe
start = time.time()
distances, indices = index.search(queries, top_k)
elapsed = time.time() - start
# Compute recall
recall_sum = 0
for i in range(n_queries):
retrieved = set(indices[i].tolist())
truth = set(gt_indices[i].tolist())
recall_sum += len(retrieved & truth) / top_k
faiss_recalls.append(recall_sum / n_queries)
faiss_times.append(elapsed / n_queries * 1000)The results confirm that FAISS reaches the same recall-nprobe relationship as our scratch implementation, but with sharply lower query times. The optimized C++ code and SIMD vectorization deliver 5-10x faster queries at the same accuracy levels, making FAISS the practical choice for production deployments.
Visualizing the Voronoi Partitioning
To build geometric intuition, let's visualize IVF in 2D.
This visualization reveals IVF's core efficiency: by comparing the query to centroids first, we identify the three most promising regions (blue, green, yellow). Only vectors in those cells are examined, while the majority of the dataset (gray points) is completely skipped. This spatial partitioning turns an search into an search, delivering large speedups.