Skip to content

feat(diskann): add configurable partitioned builds to reduce build memory - #867

Open
richyreachy wants to merge 9 commits into
alibaba:mainfrom
richyreachy:feat/diskann-partitioned-build
Open

richyreachy wants to merge 9 commits into
alibaba:mainfrom
richyreachy:feat/diskann-partitioned-build

Conversation

@richyreachy

@richyreachy richyreachy commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Add an optional partitioned build path for DiskANN to reduce peak working memory when building large indexes. Vectors are streamed to temporary storage, built into bounded subgraphs, and merged into a single disk-backed graph.

Partitioned builds are disabled by default in the public API.

Changes

  • Assign vectors to overlapping partitions and split oversized partitions into bounded chunks.
  • Merge subgraph neighbors using global vector IDs, deduplication, and robust pruning. Add a global successor cycle to preserve reachability.
  • Train one global PQ codebook and encode vectors in bounded batches using the builder thread pool.
  • Select a global entry point from the dataset mean and additional entry points from shard medoids.
  • Expose partitioned_build and build_memory_budget_gb through C++, Python, and database configuration.
  • Persist configuration across reopen operations while preserving compatibility with older manifests and Python pickle states.
  • Preserve the existing DiskANN index format and search path.

Example:

zvec.DiskAnnIndexParam(
    partitioned_build=True,
    build_memory_budget_gb=0.25,
)

Test Results

1M Cohere Dataset

  • Build
One Shot Partition Build(1 GiB Budget)
Mem Anon 3.14 GiB 0.58 GiB
RSS Peak 6.02 GiB 3.46 GiB
Build Time 2 Min 29 Sec 5 Min 12 Sec
  • Search
Build ls=100 ls=300 ls=500
One Shot 94.614 / 89.6 97.882 / 34.1 98.694 / 21.2
Part 8 GiB 94.700 / 88.7 98.006 / 33.8 98.804 / 21.0
Part 4 GiB 95.166 / 89.4 98.352 / 34.2 99.016 / 21.2
Part 2 GiB 94.590 / 91.3 98.268 / 34.5 99.008 / 21.4
Part 1 GiB 93.560 / 87.1 97.728 / 33.6 98.738 / 21.0

# Conflicts:
#	src/core/algorithm/diskann/diskann_builder.cc
#	src/core/algorithm/diskann/diskann_params.h
…ed build

The partitioned builder encoded PQ codes single-threaded, which dominated
build time (5.5 min of 9.7 min on Cohere 1M). Stripe each batch across the
builder thread pool; encode now takes ~40 s.

It also picked the medoid from the first reservoir sample and used the
farthest-point centers as entry points, which are outliers by construction.
Searches started far from the data mass and ran ~12% slower than one-shot
indexes. Accumulate the global mean while streaming the snapshot, pick the
nearest vector as medoid during shard assignment (identical to
DiskAnnBuilder::calculate_entry_point), and record each shard's medoid as an
entry point. QPS now matches one-shot builds.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant