perf(vector/index): build-time visited-set / candidate-heap reuse with thread-local arenas

[mosuka]

Round-3 perf push sub-issue (tracked under umbrella #535).

[M] Build-time visited-set / candidate-heap reuse with thread-local arenas

• Where: laurus/src/vector/index/hnsw/writer.rs (search_layer, currently around line 1243
  — the original :842-938 line reference had drifted; see investigation comment below)
  allocates HashSet visited, two BinaryHeaps, and a Vec per call).

• Current behaviour: every node insertion calls search_layer once
  per layer (so O(top_level) times per node); each call allocates a
  fresh HashSet<u64> + two BinaryHeaps. With rayon parallel build
  the allocator becomes a bottleneck.

• Why it might be a bottleneck: HashSet::new() is roughly free
  but the first inserts trigger heap allocation; the heaps grow to >= ef
  capacity. With M=16, ef_c=200, 10 M nodes, this is on the order of
  hundreds of millions of allocator calls.

• Reference precedent: hnswlib uses a thread-local "visitedList"
  pool that resets a u16 generation counter per call instead of
  clearing the set; FAISS does the same with VisitedTable.

• Current data structure: HashSet<u64> resized per call,
  BinaryHeap<Candidate> resized per call.

• Proposed layout (superseded — see investigation comment below for the adopted design, a
  BitVec + touched-list arena instead of this Vec<u32> generation-counter shape, which turned
  out to retain hundreds of MB per thread for the process lifetime):

  ThreadBuildArena {
      visited_gen : Vec<u32>    // sized to max_doc_id + 1
      gen_counter : u32         // bumped on each search; visited iff visited_gen[i] == gen_counter
      to_visit    : BinaryHeap<VisitorCandidate>   // capacity ef_c, .clear() per call
      found       : BinaryHeap<Candidate>          // capacity ef_c, .clear() per call
  }
  

  Reset is one increment + two clear()s — no allocations after warm-up.

• Suggested direction: introduce ThreadBuildArena and pass it
  through search_layer and prune_neighbors. Pair with the read-side
  arena HnswSearcher already uses for the visited-set bitmap (perf(vector/hnsw): replace HashSet visited-set with bitmap #406).
  (Superseded — see investigation comment: the read side does not actually have a thread_local!
  arena, and prune_neighbors is out of scope for this PR; see task list below.)

• Risk / scope: medium; the arena needs sizing once max_doc_id
  is known. Substantial CPU savings for large parallel builds.

Task list

• Investigation & plan (see comment below)
• Implement SearchLayerArena (BitVec + touched undo-list + two BinaryHeaps) as a
  thread_local!, wired through search_layer
• New unit tests (arena reuse doesn't leak visited state, grows correctly, error semantics
  preserved, per-thread isolation) + mutation checks
• Full verification (fmt/clippy, cargo test -p laurus, existing HNSW recall/reachability/
  append tests unchanged)
• File a follow-up issue for ConcurrentHnswGraph::get_neighbors_view's per-candidate
  Vec<u64> clone (found during investigation — likely a bigger allocation source than this
  issue's original three containers, but out of scope here) — perf(vector/index): ConcurrentHnswGraph::get_neighbors_view clones a fresh Vec<u64> per call on the build-time hot path #1137
• PR opened — perf(vector/index): reuse search_layer's visited-set and candidate-heaps via a thread-local arena #1138

---

---

ID: VI-14 — see ~/.claude/tasks/laurus/20260523_perf_round3_audit/task_list.md for the full Round-3 issue list.

[mosuka]

Investigation

Read search_layer, its sole call site (insert_one), prune_neighbors, calc_dist, and the
doc_id_map/vectors data flow directly, then had an independent design review pass on the
proposed arena.

Two things in this issue's description don't match current code:

• The file:line reference (writer.rs:842-938) is stale — the logic itself is unchanged, it's
  just moved (now around line 1243) as the file grew from other commits.
• "Pair with the read-side arena HnswSearcher already uses for the visited-set bitmap (perf(vector/hnsw): replace HashSet visited-set with bitmap #406)" is
  inaccurate: hnsw/searcher.rs allocates a fresh BitVec::from_elem(...) on every search call —
  there's no thread_local! there. The only thread_local! in the whole crate is a test-only RNG
  helper (lexical/index/structures/norms.rs). The read side's dense-ordinal-indexed bitmap
  convention is still useful precedent, just not an existing reusable arena to pair with.

Decision: BitVec + touched-list instead of the proposed generation-counter Vec

The proposed visited_gen: Vec<u32> design has a real problem: it's retained per OS thread for
the process lifetime (rayon's global pool threads never die), so at 10M nodes it holds ~40 MB per
thread — roughly 640 MB across a 16-thread pool — permanently, versus the HashSet it replaces,
which only ever held the few hundred actually-visited ids and was freed after each call.

Switching to a BitVec (1 bit/node, same convention the read-side searcher already uses) plus a
touched: Vec<usize> undo-list that records only the indices set during the current call brings
this down to ~1.25 MB/thread at 10M nodes, and the reset cost becomes proportional to what was
actually visited (typically bounded near ef_construction) rather than the whole node count. It
also eliminates the generation-counter wraparound handling entirely, since there's no counter left
to overflow.

Scope decisions

• prune_neighbors is left out of this PR. It doesn't actually use a HashSet/BinaryHeap at
  all (just a small Vec<Candidate>, sorted and truncated, only when over max_conn — bounded by
  m/m_max_0, ~16-32 items). The issue's own proposed struct doesn't include a field for it
  either. It's also called from ConcurrentHnswGraph::add_neighbor_with_pruning while holding a
  node's write lock, so sharing a thread-local arena with it adds lock-interaction risk for little
  benefit.
• ConcurrentHnswGraph::get_neighbors_view's lock.read().clone() turns out to likely be a
  bigger allocation source than the three containers this issue names — it allocates a fresh
  Vec<u64> every time search_layer pops a candidate (and again during the Phase A greedy
  descent). Fixing it means changing the GraphView trait's signature to write into a
  caller-provided buffer, touching both HnswGraph and ConcurrentHnswGraph — a separable, larger
  change. I'll file a follow-up issue for it once this PR lands.

Plan

• search_layer's HashSet<u64> + two BinaryHeaps become a thread_local! SearchLayerArena
  (BitVec + touched: Vec<usize> + the two heaps), reset on entry (not exit, so an
  error/early-return from the previous call never leaves stale state for the next one).
• New calc_dist_by_idx helper (extracted from calc_dist) lets the hot loop resolve a neighbor's
  dense index once and reuse it for both the visited-check and the distance calculation, avoiding a
  double doc_id_map lookup for newly-visited neighbors — this keeps the already-visited case at
  one lookup (previously zero) but the newly-visited case at one lookup (previously two via
  visited.insert + calc_dist's own lookup), so it doesn't regress net lookup count.
• search_layer's return type changes from Result<BinaryHeap<Candidate>> to
  Result<Vec<Candidate>> (no need to re-heapify on the way out); select_neighbors's parameter
  type changes from &BinaryHeap<Candidate> to &[Candidate] accordingly (body unchanged, and the
  one caller needs no changes — &Vec<Candidate> auto-derefs).
• New unit tests live in a new #[cfg(test)] mod tests inside writer.rs itself (the existing
  hnsw::tests module is a sibling of hnsw::writer, declared in hnsw.rs, so it can't reach
  the private items this change touches): arena reuse doesn't leak visited state across calls, the
  BitVec grows correctly when the vector set grows between calls, a neighbor missing from
  doc_id_map still errors (proving the lookup-reordering didn't swallow that case), and the arena
  is genuinely per-thread. Each gets a mutation check (temporarily break the relevant reset/lookup
  logic, confirm the test goes red, then restore).
• No functional/behavioral change to search results — the existing HNSW recall/reachability/append
  test suites are the regression net for that.

No breaking changes; internal implementation only, no public API surface touched.

Task list tracked in the issue description.

[mosuka]

Filed #1137 for a bigger allocation source found while implementing this: ConcurrentHnswGraph::get_neighbors_view's per-pop Vec<u64> clone. Out of scope for this PR (see the investigation comment above and #1137 for why).