feat(core): Add new index PartitionedHNSW for vectors - #9469
Open
ghost wants to merge 17 commits into
Open
Conversation
ghost
self-requested a review
July 16, 2025 02:38
github-actions
Bot
added
area/schema
github-actions
Bot
added
the
area/integrations
shivaji-kharse
force-pushed
the
harshil-goel/split-vector3
branch
from
August 29, 2025 06:47
1978bda to
f2cade4
Compare
|
This PR has had no activity for 60 days and has been marked stale. Comment to keep it active. |
shiva-istari
force-pushed
the
harshil-goel/split-vector3
branch
from
August 4, 2026 17:50
037faa1 to
dd76371
Compare
This comment has been minimized.
This comment has been minimized.
- MergeResults: stop at len(result) instead of indexing past it when there are fewer candidates than maxResults. - kmeans updateCentroids: keep the previous centroid when a cluster receives no vectors in a pass instead of dividing by zero into NaNs; drop debug prints from the hot loop. - partitionedHNSW.Insert: learn the vector dimension from the first vector when it is unset (was a self-assign no-op that rejected every insert on a fresh predicate). - Reject partitionStrat "query": it passed validation but left the partition nil and panicked on first use. Only kmeans is implemented. - rebuildVectorIndex: propagate BuildInsert errors instead of dropping them; route pass logging through glog.V(1).
- CreateKMeans takes numClusters and numProbes; NumSeedVectors returns numClusters instead of a hardcoded 1000, so non-default cluster counts no longer route inserts past the cluster map. - New numProbes schema option (default max(4, numClusters/25), clamped to [1, numClusters]) — how many clusters a search will visit once routing is wired in. - applyOptions validates numClusters >= 1. - SetNumPasses(0) clears the seed centroids so a degenerate build (fewer vectors than clusters) persists an empty centroid set and consistently routes through cluster 0. - vectorCentroids gets an RWMutex: build passes write, routing reads — needed once the index instance becomes long-lived and serves concurrent lookups.
…ed hnsw The centroid router existed (findNClosestCentroids) but nothing called it: FindIndexForSearch ignored the query vector and returned every cluster, so each similar_to fanned out to all shards, and persisted centroids were never loaded, so any fresh index instance routed all inserts to cluster 0. - VectorPartitionStrat.FindIndexForSearch/FindIndexForInsert now take an index.CacheType so the strategy can lazily hydrate persisted centroids on first use (nil during builds — the build owns the centroids in memory). - kmeans hydrates once per instance, caches a miss (never-built predicate = consistent cluster-0 mode), and routes searches through findNClosestCentroids honoring numProbes. - partitionedHNSW.Search fan-out now runs through a bounded errgroup (2*GOMAXPROCS) and propagates shard errors; clusters that are merely empty contribute zero results instead of failing the query. - SearchWithUid implemented (was a cluster-0 stub): fetch the uid's vector via the new hnsw.GetVectorFromUid, route like Search, drop the query uid when the filter demands it. SearchWithPath searches the vector's own cluster instead of an arbitrary one. - partitionedHNSW implements OptionalSearchOptions, so per-query ef and distance_threshold reach the shards instead of being silently dropped.
Every mutation and query used to call FactoryCreateSpec.CreateIndex, which maps to CreateOrReplace: a brand-new partitionedHNSW (with empty routing centroids and 1000 fresh sub-indices) was constructed per operation, so live inserts always routed to cluster 0 and searches re-created the world on every call. - IndexFactory gains FindOrCreate. The partitioned factory returns the existing instance (preserving hydrated centroids); the plain hnsw factory keeps fresh-instance semantics (persistentHNSW's nodeAllEdges/deadNodes maps are transaction-scoped caches, only safe per-call). - posting.addIndexMutations and the similar_to query path use the new FactoryCreateSpec.FindOrCreateIndex; rebuildVectorIndex keeps CreateIndex, so a rebuild atomically swaps in the freshly built instance. - partitionedHNSW hands out a fresh persistentHNSW view per operation for just the probed clusters (subIndex) instead of sharing the build-time clusterMap: sharing those would race on their per-instance caches and leak edge state across transactions. Net effect is still far cheaper than before — an operation now creates at most numProbes sub-index views, not numClusters.
A vfloat DEL always appended the uid to the unsplit <pred>__vector_dead list, but partitioned sub-indices only consult their own split <pred>__vector_dead_<i> lists — so deleted vectors kept coming back in partitioned search results. New optional index.VectorDeadListResolver interface: partitionedHNSW routes the deleted vector like an insert and returns that cluster's dead attr; posting's DEL branch uses it when the index implements the interface. Plain hnsw behavior is unchanged. Split-attr naming is now exported from tok/hnsw (SplitEntryAttr/SplitVecAttr/SplitDeadAttr) so delete routing and the upcoming rebuild cleanup share one definition.
…t changes - prefixesToDropVectorIndexEdges only dropped the unsplit __vector_* attrs. Partitioned per-cluster attrs (pred__vector_entry_<i> etc.) are distinct length-prefixed predicates, so a reindex or index drop left every cluster's data and the persisted centroids behind. Now enumerates the per-cluster prefixes for every cluster count found in the old or current schema spec (covering numClusters changes in both directions) plus the centroid key. - Factory identity: numClusters and partitionStratOpt now participate in the factory spec name, so a numClusters change is detected as a rebuild. vectorDimension stays excluded (SetDimension auto-appends it to the stored schema; including it would make every re-apply look like a change), as does the query-time-only numProbes. - Seed selection: the rebuild reservoir-samples numClusters seeds across the full scan (seeded by StartTs for retry determinism) instead of taking the first N vectors in badger key order, which biased the initial centroids. - A degenerate build (fewer vectors than clusters) needs no special persist step anymore: the pre-rebuild drop removes any stale centroid key and the hydration miss keeps routing in cluster-0 mode.
New TestPartitionedPipelines drives the four supported pipelines on a real cluster: index build over existing data, query routing, live inserts after the build (must find themselves via similar_to), delete-then-search (the deleted uid must disappear), an alpha restart (search and insert routing must re-hydrate centroids from disk), and a numClusters change (rebuild to a different layout keeps every vector findable).
…on test The per-cluster graph commit in rebuildVectorIndex ran through x.ExponentialRetry(int(x.Config.MaxRetries), ...) with the error ignored. With max-retries unset (any process that is not a fully configured alpha) that is zero attempts: the commit silently never executed and the entire cluster graph was lost. The retry now makes at least one attempt and its error is checked, and the TxnWriter is flushed so async commit errors stop being discarded (same for the centroid persist). Centroid hydration logs its outcome at v=1 so a restarted alpha's routing state is observable in logs. New posting/vector_restart_test.go covers the restart contract deterministically: build a partitioned index, then verify a brand-new index instance (what a restarted process has) hydrates persisted centroids and routes every vector to itself — including after forced rollups of all aux keys and after a full replayed rebuild (drop + rebuild at the same StartTs).
The restart subtest is inherently flaky for reasons outside the partitioned index: an alpha restart replays the schema alter from the raft WAL, which drops and re-runs the full index rebuild asynchronously while the replayed data mutations race it — a mutation routed by mid-training centroids (or wiped by the replay's DropPrefix) becomes an unreachable graph node until the next rebuild. Pre-existing reindex-vs-mutation race, affects all index types, needs the mutation-pipeline serialization work. Readiness polls, pre-restart raft snapshots and waiting out the opIndexing task were all tried and cannot close the window from the client side. Restart hydration itself is covered deterministically by posting/vector_restart_test.go.
partionedhnsw is no longer a separate user-facing index type. The hnsw tokenizer now dispatches on the numClusters option: absent keeps today's monolithic index byte-for-byte (including its factory identity string, so existing predicates do not re-index on upgrade); numClusters > 1 engages the partitioned implementation. Partitioned-only options (numProbes, partitionStratOpt, vectorDimension) without numClusters are rejected with a clear error instead of being silently ignored. Spec recognition everywhere (SetDimension, rebuild drop-prefixes, backup) switches from matching the index name to checking numClusters presence via partitioned_hnsw.SpecHasOption. Experimental guards: the bulk loader and predicate move reject partitioned specs with actionable errors; export now skips the persisted centroid keys (also fixes the centroid-leak-on-export bug). Tests: unified factory dispatch/identity/flip-transition, the plain-hnsw back-compat no-reindex pin, and all existing partitioned tests updated to the hnsw(numClusters:...) schema syntax.
TestVectorIndexDropPredicate, TestVectorIndexWithoutSchema and TestIndexRebuildingWithoutSchema asserted similar_to returns exactly topK results. That holds for monolithic hnsw but not for a partitioned index, where the result count is bounded by the probed clusters' contents: these tests use numClusters == numVectors (~1 vector per cluster), so a default numProbes (numClusters/25 = 40) yields ~40 results, not 100. The same assumption sat in the index-readiness Eventually() polls, which then timed out for partitioned. Assert index functionality (non-empty, every result a real inserted vector) for partitioned while keeping the exact-topK check for monolithic. Pre-existing on the branch; surfaced by CI.
The shared partitioned test schema used numClusters=1000 against ~1000-vector datasets — ~1 vector per cluster, which defeats clustering and, with the default numProbes (numClusters/25=40), structurally caps similar_to at ~40 results. Tests asserting topK=100 results then failed for the partitioned iteration (a config problem, not an index bug: all vectors store/restore fine; the search just never probes clusters it wasn't told to). Set numClusters=8 (~125 vectors/cluster) so a default numProbes gathers well over topK candidates and the existing strong assertions hold for both index types. Supersedes the assertion-weakening in the previous commit. Dedicated pruning tests (TestPartitionedPipelines, TestPartitionedHNSWIndex) keep their own inline schemas and are unaffected.
Enable bulk loader to build partitioned (IVF-over-HNSW) vector indexes by deferring the build to the post-reduce phase. Raw vectors are streamed to the shared tmpDb during reduce, then RebuildVectorIndexForBulk runs the multi-pass build (10 passes, ~numClusters/10 graphs per pass) to bound peak memory while preserving maximum speed. Key changes: - Add SkipVFloatConversion field to IndexRebuild; skip pre-pass for bulk (fix: same-key-same-ts overwrite hazard where re-writing at writeTs collides with reduce output) - Fix RunWithoutTemp tail: guard ExponentialRetry(MaxRetries) to ensure at least one attempt (fix: silent skip when MaxRetries=0 outside alpha) - Add RebuildVectorIndexForBulk wrapper to drive alpha's rebuild machinery - Add NumClustersFromSpec helper to unified_factory for single source of truth - Replace streaming insert with deferred build for partitioned predicates: modify toList to track vecNone/vecStreaming/vecDeferred; skip append for deferred; call trackPredShard on first classification - Add isDeferredPred and trackPredShard methods to vector_indexer - Add buildDeferredVectorIndexes method to run sequential per-predicate builds - Extend copyVectorDataToShards to handle per-cluster split attrs and centroid Tests: - TestNumClustersFromSpec: verify helper extracts cluster counts correctly - Existing vector_restart_test already covers programming patterns Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
Optimize the multi-pass partitioned index build with two complementary improvements that remove redundant work: 1. Routing memo (tok/partitioned_hnsw): During the 10 index passes, centroids are frozen. BuildInsert computes nearest-centroid routing (O(numClusters·dim)) on EVERY vector in EVERY pass. Memoize uuid→cluster via sync.Map during index passes (capped at 32M entries to bound memory); subsequent passes reuse the cached routing. Avoids 10·N·k·d redundant mult-adds (~10^15 at N=1M, k=1000, d=100+). Safety: centroids locked during passes; racing double-computes store identical values (benign); sync.Map fits write-once/read-many profile. Benefits both bulk builds and alpha alters. 2. Parallel per-cluster commits (posting/index.go): Each cluster is independent (disjoint keyspace via UpdateIndexSplit); replace the sequential commit loop with an errgroup (SetLimit to min(finished, GOMAXPROCS)). Each goroutine owns its TxnWriter and ExponentialRetry; badger.CommitAt is concurrent-safe. Parallelizes write I/O and raft proposal latency. Testing: existing suite verifies correctness; -race flag guards memo and errgroup against concurrency bugs. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
Optimize k-means training for bulk-loaded partitioned vector indexes by replacing full disk scans with in-memory sampling. Industry practice (FAISS) trains coarse quantizers on ~39-256 points per centroid; more adds negligible quality. Sampling removes 5 full IterateDisk passes plus (1−S/N) of the k·d assignment flops (~40× less k-means work at N=10M, k=1000, d=100+). Sample size S = min(N, 256·k) is memory-resident (S·d·4B ~131MB at d=128); capped at 1GB by shrinking S with floor max(numSeeds, 32·k). Deterministic: reservoir seeded with StartTs. Changes: - Add SampledKMeans bool to IndexRebuild; bulk wrapper sets it true. - Extend seed reservoir scan to support variable sample size (same single IterateDisk pass). - Sampled k-means path: replace full disk scan with in-memory sweep of sampled vectors, parallelized via errgroup across runtime.GOMAXPROCS(0). - Micro-tuning: adapt stream.NumGo to max(16, GOMAXPROCS) for better multicore utilization. - Free sampled vectors after k-means to recover memory before index passes. Integration tests: add TestBulkLoadPartitionedVectorIndex to systest/vector. Testing: existing suite verifies correctness; sampled vs full k-means produce equivalent centroids (random seed ensures determinism). Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
The sampled-k-means training-sample size clamped the 256*k target to the 1GB RAM budget via min(estimatedSize, maxByRAM), but the very next line unconditionally overwrote it with max(numSeeds, max(32*k, 1)) — making the RAM cap and the 256*k target dead code. The sample was always 32*k. Fold the floor into the clamp so the size is max(numSeeds, max(32*k, min(256*k, RAM-cap))): targets 256 vectors per centroid, capped by the RAM budget, floored so a tiny-dimension cap can't starve training. At k=1000,d=1536 this is ~174,762 (was 32,000). Verified on dbpedia-openai-1M: recall@10 unchanged (0.92), and the RAM cap now actually bounds the training sample as documented. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
shiva-istari
force-pushed
the
harshil-goel/split-vector3
branch
from
August 13, 2026 07:13
68da927 to
46bac2a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.