Partitioned HNSW for large vector corpora: design and benchmarks (up to 1M×960d)

[shiva-istari]

We've been working on the partitioned HNSW vector index (PR #9469 lineage) — getting it correct and benchmarking it honestly across three dataset scales. This post shares the design and the numbers (including a 1M-vector run with exact ground truth). Feedback welcome, especially on the scale trade-offs.

⚠️ Work in progress / experimental. This is an active development effort, not a released feature. The results below are from a development branch; the API, defaults, and supported surface are all subject to change. Do not use in production.

What it is (in one paragraph)

Plain HNSW builds one large proximity graph over all vectors. Partitioned HNSW is IVF-over-HNSW: k-means splits the vectors into numClusters clusters, builds one small HNSW graph per cluster, and persists the cluster centroids. A query compares against the centroids, then searches only the nearest numProbes clusters and merges the results. It is not a reimplementation — each cluster is a real HNSW graph built by the existing engine; partitioning is a ~1000-line routing/orchestration layer on top.

What works today

Four core pipelines are implemented and covered by unit + integration tests:

• ✅ Index build (schema alter over existing data — the multi-pass rebuild pipeline)
• ✅ Index rebuild (including detecting a numClusters change and re-hydrating routing state after an alpha restart)
• ✅ Mutations (live inserts route to the correct cluster; deletes route to per-cluster dead lists)
• ✅ Queries (similar_to centroid routing, cross-cluster fan-out and merge; per-query ef/distance-threshold options)

Everything outside those four is not supported yet (see Known limitations). The intent is to land the correct core first and add the remaining surface incrementally.

Benchmark setup

• Build: branch harshil-goel/split-vector3, single zero + alpha via docker compose, 24 GB machine (Docker VM limit 23.4 GB).
• Index-build memory measured on the rebuild path — load data without an index, then create the index via schema alter (the multi-pass build pipeline). Alpha restarted before each build for a clean baseline; Go heap (go_memstats_heap_inuse_bytes) + RSS sampled every 250 ms; build time = alter → opIndexing cleared.
• Identical index params both configs: maxLevels:3, efSearch:32, efConstruction:64, metric:euclidean. Partitioned adds numClusters ≈ √N (400 at 171k, 1000 at 1M); numProbes defaults to numClusters/25 (16, 40).

dataset	vectors	dims	source	ground truth
products	3,356	384	dgraph-benchmarks/vector suite	exact KNN (sklearn)
TREC-COVID (BEIR)	171,332	768	real text embeddings (all-mpnet-base-v2)	BEIR relevance judgments (NDCG)
GIST1M (ann-benchmarks)	1,000,000	960	real image descriptors	exact top-100 neighbors shipped with the dataset

Result 1 — index-build memory (the headline)

Peak heap above the post-restart baseline:

dataset	monolithic	partitioned	saving
171k × 768d	+3.85 GiB	+3.40 GiB	−12%
1M × 960d	+21.9 GiB	+11.0 GiB	−50%

At 1M the monolithic build peaked at 23.4 GiB — flush against the machine's ceiling — while the partitioned build peaked at 12.5 GiB in a steady per-pass sawtooth (build a slice of clusters → commit → free → next).

Why the saving grows with scale. Build memory ≈ a roughly-flat machinery floor (badger memtables, caches, GC churn — paid by both) + the held graph state (edges + uncommitted deltas) that partitioning splits across passes. At 171k the held state is ~0.5 GiB — splitting it barely moves the total. At 1M it's ~17+ GiB and dominant, and it grows super-linearly. One corpus-size step further, the monolithic build simply won't fit on ordinary hardware while the partitioned peak stays bounded by pass size. The −12% at small scale and −50% at 1M are the same mechanism at different points on that curve. (171k reproduced within ~0.05 GiB across two runs.)

Result 2 — accuracy

dataset	metric	monolithic	partitioned
TREC-COVID 171k	NDCG@10 (BEIR qrels)	0.459	0.477–0.481 (two runs)
GIST1M	recall@10 vs exact ground truth	0.286	0.753 (2.6×)

At identical ef params, partitioned was more accurate — and dramatically so on GIST (the classic hard case). One greedy walk through a 1M-node graph at efSearch=32 collapses to 0.29 recall; 40 independent searches over ~1k-node cluster graphs explore far more of the true neighborhood and hold 0.75. Routing correctness was verified separately: every indexed vector returns itself as top-1 (self-recall = 1.00), zero errors across ~1.5M vectors indexed.

Result 3 — build time

dataset	monolithic	partitioned
171k	2.4 min	2.8 min
1M	153.6 min	~140 min

Par at both scales. (Side-finding: the rebuild path is ~11× faster than inserting into a live index — 2.5 min vs 29 min at 171k, either config. "Load data, then add the index" is the cheaper bulk-ingestion strategy.)

Result 4 — query latency (the trade-off)

dataset	monolithic	partitioned
171k (16 probes)	p50 116 ms	p50 112–139 ms
1M (40 probes)	p50 232 ms / p95 2.3 s	p50 1.3 s / p95 5.2 s

numProbes is a latency↔recall dial. At 171k with moderate probes, partitioned matched monolithic latency and beat its NDCG. At 1M the 40-probe setting bought the 2.6× recall at a latency cost. Note a same-recall comparison would force the monolithic efSearch much higher — slower and hungrier at build time — so at matched params partitioned wins memory and accuracy and pays in latency.

Where partitioning does NOT help

• Live-insert path (index present while inserting): memory is parity — it's dominated by shared mutation machinery, which partitioning has no lever on. The build-memory benefit is a rebuild-path property; docs should say so.
• Small corpora: below ~100k vectors the machinery floor dominates and the win is marginal. This is a large-corpus feature.

Correctness work that made this measurable

The partitioned index required a fix stack before it was functional: wiring the (previously dead) centroid search-routing, persisting + re-hydrating centroids across restarts, routing deletes to per-cluster dead lists, rebuild data hygiene (dropping stale per-cluster keys, detecting numClusters changes), and a build-commit durability fix (a retry bug silently dropped cluster graphs in non-alpha processes). Unit + integration coverage now spans all four supported pipelines (build, rebuild, mutation, query) from a starting point of 0%.

Known limitations / not yet supported

• Lifecycle ops for partitioned specs: bulk loader, export, predicate move, incremental backup / restore-then-reindex. Bulk loader and predicate move now fail fast with a clear message; export skips the centroid key (fixing a leak); the rest are follow-ups.
• Restart-during-writes: a pre-existing reindex-vs-mutation replay race (affects all index types) can leave a few nodes unreachable until the next rebuild. Needs the mutation-pipeline serialization work.

Try it out

This is early, experimental work and we'd love more eyes on it. Play with it on your own datasets, and let us know how it behaves — feedback, suggestions, and contributions are all very welcome.

[shiva-istari]

Update: the bulk loader now builds partitioned indexes

One of the "not yet supported" items above is done: dgraph bulk can now build a hnsw(numClusters:N) predicate offline. You no longer have to load first and add the index via a live schema alter — the partitioned index comes out of the bulk pipeline directly.

How it works

Normally dgraph bulk reads your data, converts it to Dgraph's internal key-value form, and builds the index as it goes. That doesn't work for a partitioned vector index: the cluster centers come from k-means, and k-means needs to see the vectors first — until the centers exist, every vector would fall into one cluster.

So the build is split into two stages:

1. During the normal map/reduce, just write the raw vectors to the temp store — no index yet.
2. After all vectors are written, run one extra pass that trains the k-means centers on a sample of the data, then routes every vector to its nearest center and inserts it into that center's small per-cluster graph. The finished cluster graphs and centers are copied into the output.

The result is a ready-to-serve partitioned index straight out of dgraph bulk — no live re-index step afterward. It reuses the same graph-builder the schema-alter path uses; only the "when" changed.

Three things keep that second pass cheap: sampled k-means (centers trained on a bounded, RAM-capped sample rather than repeated full-corpus scans), a routing memo (each vector's cluster is computed once and reused across the index passes instead of recomputed every pass), and parallel per-cluster commits.

Results

Two datasets, monolithic vs partitioned, identical index parameters. Recall@10 is measured against exact euclidean brute-force ground truth.

dataset	metric	monolithic hnsw	partitioned hnsw(numClusters)
GIST — 200k × 960d (k=448)	build time	7.5 min	1.8 min (4.3×)
	peak RSS	6.70 GiB	6.12 GiB
	recall@10	0.386	0.778
dbpedia-openai — 1M × 1536d (k=1000)	build time	84.1 min	21.8 min (3.9×)
	peak RSS	14.38 GiB	13.3 GiB
	recall@10	0.799	0.921

dbpedia-openai = 1M real OpenAI text-embedding-ada-002 embeddings (1536d).

A note on the build-time gap: these numbers are at badger's default settings. Part of the monolithic slowdown is I/O — the logs show the 1M-node graph doesn't fit the default block cache (hit-ratio: 0.00), so lookups thrash to disk and the insert rate decays from ~5000/sec to ~217/sec. Tuning badger (a bigger block cache) would help the monolithic case; we haven't measured by how much.

But tuning wouldn't erase the gap, because the real reason is structural: a monolithic insert has to search a graph of all N nodes, and that search only gets more expensive as the corpus grows. A partitioned insert only ever touches its own ~N/k-node cluster graph (~1000 nodes at k=1000), which stays small and cheap no matter how big the corpus gets. So partitioning's build advantage is inherent to the data structure, not just a caching effect.

Higher recall is the usual IVF effect — probing ~40 small cluster graphs covers more of the true neighbourhood than one greedy walk at efSearch:32, which is also why partitioned query latency is higher. On the bulk path peak memory is roughly a wash (unlike the −50% on the rebuild path), since the monolithic bulk build already commits-and-frees incrementally; here the win is build speed and recall.

Exactly how it was run

• Binary / machine: native darwin/arm64 build of branch harshil-goel/split-vector3, CGO_ENABLED=0 (Go allocator, no jemalloc). macOS, 24 GB RAM, 12 cores. Build time = bulk wall-clock; peak memory = max RSS from /usr/bin/time -l.
• Data → RDF: one N-Quad per vector, UID = row index + 1:

  <0x1> <embedding> "[0.017399,-0.014088, ...]" .
  

  GIST from the ann-benchmarks HDF5 (one RDF file); dbpedia from the HF parquet (26 gzipped RDF shards).
• Schema — identical apart from numClusters:

  # monolithic
  embedding: float32vector @index(hnsw(maxLevels:"3", efConstruction:"64", efSearch:"32", metric:"euclidean")) .
  # partitioned  (numClusters ≈ √N: 448 at 200k, 1000 at 1M)
  embedding: float32vector @index(hnsw(maxLevels:"3", efConstruction:"64", efSearch:"32", metric:"euclidean", numClusters:"1000")) .
  

• Bulk loader — single output group, default badger:

  dgraph bulk -f <rdf> -s <schema> --zero localhost:5080 \
              --map_shards 1 --reduce_shards 1 -j 8 --out ./out
  

  --reduce_shards 1 → one group / one p directory; -j 8 worker threads; badger left entirely at defaults (block cache, memtables — nothing tuned).
• Serve & verify: a fresh zero + alpha -p out/0/p on the bulk output — confirmed the schema persisted (hnsw(...numClusters:"1000")) and count(has(embedding)) == 1,000,000, and spot-checked that each query vector returns itself as its own top-1.
• Query & recall: similar_to(embedding, 10, <vec>); partitioned uses the default numProbes = numClusters/25 (≈18 / 40). Recall@10 = overlap of the index's top-10 with exact brute-force top-10, averaged over 200 queries.

These are single runs per config at 1M on a development branch, so treat the exact multipliers as one operating point rather than a tuned frontier — numClusters, numProbes, efSearch and the badger cache size all move the trade-off.

Still not supported for partitioned specs

Export, predicate move, and incremental backup — those remain follow-ups. Feedback welcome, especially from anyone bulk-loading large real-world embedding corpora.