You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
diskgraph: record sampled edges through per-seed maps - #101
This PR makes multi-hop DiskGraph sampling 1.3–2× faster per thread by changing how each layer records sampled edges. The output is identical to the current path.
Problem. Deep sampling plans spend most of their time after the edges are read:
Why does 5-hop sampling on a 1B-row graph run at ~330 seeds/s per thread when reading the adjacency takes a fifth of the time?
Why do 8 sampler threads give only ~1.3× one thread at 5–6 hops (about 1.9 cores busy)?
Cause. Each layer runs in two phases. The map phase reads adjacency in parallel, one shard per task. The commit phase then runs serially: for every sampled edge, record_edge looks up (seed, node) in one index_map per node type, which holds every node sampled so far for every seed in the request. At b=512 and 5 hops that map reaches millions of entries, so nearly every lookup misses cache. Profiling 5-hop sampling on one thread: commit (record_edge and its hash-map lookups, inserts, rehashes and page faults) took about 70% of the CPU and the parallel map phase about 16%.
Change.commit_edges replaces the per-edge loop for the standalone (non-fused) path:
Bucket the layer's edges by seed (a counting sort), carrying each edge's node key and position.
For each seed, deduplicate against that seed's own small map (seed_maps[slot], about 10k nodes at 6 hops, which stays in cache). Each edge resolves to an existing local id, a new node, or a repeat of an earlier edge in this layer.
Walk the edges in their original order: assign local ids to new nodes, then emit every edge. Ids, edge order, frontiers and stats are exactly what record_edge produced.
Write the new ids back into the per-seed maps.
The per-shard map-phase edge buffers are also kept in SamplingState and reused across node types, layers and requests instead of being reallocated. DiskGraphEngine::legacy_recording = true selects the previous path; the fused path is unchanged.
Back-compat
Situation
Behaviour
Standalone sampling (fused=false)
Same local ids, edges, frontiers and stats as the previous path (tested, 2–6 hops).
Fused feature-fetching sampling
Unchanged (previous path).
legacy_recording = true
Previous path.
More than 256 node types
Falls back to the previous path (the per-seed key packs the type into 8 bits).
Pooled state between requests
Per-seed maps are cleared and kept only while their summed capacity fits the existing RESET_THRESHOLD (262k entries, about 4 MiB); the rest are freed. Edge buffers and commit scratch are cleared and shrunk to RESET_THRESHOLD, as index_map already is. One heavy request does not pin its peak memory in the pool.
More than 2^31 nodes of one type in one request
Returns an error (the commit tags pending entries with bit 31); the previous path allowed up to 2^32.
TypeState::index_map
Holds only the seeds on this path; nothing outside the recording step reads it.
Measured results
Indexes: synthetic relational graphs with 7 node types (2 entity tables and 5 event tables), 106.5M rows (5.3 GB index, 5 shards per type) and 1.01B rows (50 GB, 16 shards per type). Workload: temporal most-recent sampling with fanouts 24, 8, 6, 8, 6, 8 (first h used), 5,000 seeds sharing one anchor time, b=512 per request, warm page cache. Measured with a kumo-stress-tests ctxbench-style driver that calls DiskGraphEngine::sample in-process (and fetch_features_for_type for the feature-fetch columns). Each cell is the better of two interleaved rounds. Machine: 64 vCPU, 249 GiB RAM, shared with other jobs, so expect about ±10%.
Training path (sample_runtime_batch, link prediction)
The table above uses pooled DiskGraphEngine::sample. Training calls sample_runtime_batch, which builds a fresh SamplingState per batch, so no maps or buffers carry over between batches. This run goes through that call on the 1B index with a two-tower link-prediction stage table (200k rows of customer → merchant purchases, entity = customers, target = merchants, one engine negative per row), so each batch seeds three groups: 512 customers, 512 positive merchants and 512 negative merchants. Same fanouts, temporal most-recent; features skipped through the feature projection unless noted. Per-batch time, better of three interleaved rounds of 4 batches, same machine (±10%). legacy_recording = true is the previous path in the same binary; both paths returned identical node counts in every cell.
threads
2 hops
3 hops
4 hops
5 hops
6 hops
2
89 → 84 ms (1.05×)
745 → 549 ms (1.36×)
1.72 → 1.32 s (1.30×)
7.23 → 4.36 s (1.66×)
16.4 → 10.3 s (1.59×)
8
73 → 67 ms (1.08×)
651 → 459 ms (1.42×)
1.47 → 1.05 s (1.40×)
6.63 → 3.73 s (1.78×)
14.5 → 8.19 s (1.77×)
At 2 hops this PR alone is 10–14% slower at 1 thread (31,250 vs 36,496 on 100M; 27,473 vs 30,675 on 1B): the commit phase is small there, so the extra bucketing pass is not paid back. #102 recovers it (0.96–0.99× of main). Peak RSS rises by 0.0–0.3 GB.
Known limitations
The commit phase is still serial; with 8 threads it is now most of the time at 5–6 hops. The next PR in the stack parallelizes the per-seed step.
Gains come from deduplication within each seed, so they shrink when each seed's neighbourhood has few repeated nodes.
Testing Done
New per_seed_recording_matches_legacy_recording: a random 41-node graph with repeated neighbours, temporal and uniform modes, both edge directions, repeated seeds, 1–700 seeds per request at 4 threads, and pooled state reused across requests; edges, nodes and stats equal the previous path.
New per_seed_recording_matches_legacy_recording_across_seed_types: two node types (A→B, B→A) with seeds of both types at overlapping indices, 3/40/700 seeds, 1 and 4 threads; covers the per-type slot offsets and the type bits of the per-seed key.
New tests/per_seed_recording_real_graph.rs (#[ignore], runs with DISKGRAPH_VERIFY_INDEX=<index>): 6 batches of 512 seeds, alternating a shared and per-seed anchors, 2–6 hops; checks this path equals the previous one and that the previous one is identical across fresh states. Passed on both indexes above (up to 14.4M edges per hop count).
New reset_caps_edge_buffers_and_commit_scratch: after a request that grew the edge buffers, commit scratch and per-seed maps past RESET_THRESHOLD, the next acquired state holds at most that capacity in each.
cargo test -p diskgraph-engine, -p diskgraph-ingest --test seed_table_ingest --test resolution_runtime_batch and -p diskgraph-pybind pass; cargo clippy --workspace --all-targets -- -D warnings is clean.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Walkthrough
DiskGraphEngine adds a legacy_recording flag and selects per-seed recording for eligible non-fused requests with at most 256 node types. SamplingState stores per-seed maps, commit scratch, and reusable raw edge buffers. The batch path commits edges and shares edge emission with sequential recording. Unit and ignored real-graph tests compare per-seed results with legacy results.
The previously identified pooled-memory risk has been addressed for the demonstrated shard counts. No remaining issue warrants delaying the merge after normal checks.
Pre-merge checks: 2 · 1
Failed checks: 1 warning
Check name
Status
Explanation
Resolution
Change Map Design Review
The prior compatibility defect remains. The PR adds node_key(type_id, global_id) as (global_id << 8) | type_id with only a debug_assert!(global_id < 1 << 56) (engine.rs:472-476). In release bu…
Use a collision-free key that retains the full u64 ID and type, or disable per-seed recording when any reachable ID is at least 2^56. Document the compatibility condition in the change map and add a release-compatible regression test wi…
Passed checks: 2 passed
Check name
Status
Explanation
Pr Size
The pull request changes 4 files. The limit is 40 files, so the PR size check passes.
Public Resource Access
The pull request adds no URL, repository location, or package source location for an external resource. The new test documentation uses local Cargo identifiers and a generic <index dir> environment-…
Full details: Public Resource Access
Explanation
The pull request adds no URL, repository location, or package source location for an external resource. The new test documentation uses local Cargo identifiers and a generic <index dir> environment-variable placeholder, not a public resource reference.
Full details: Change Map Design Review
Explanation
The prior compatibility defect remains. The PR adds node_key(type_id, global_id) as (global_id << 8) | type_id with only a debug_assert!(global_id < 1 << 56) (engine.rs:472-476). In release builds, valid u64 IDs such as 0 and 2^56 collide. The new path is enabled for every non-fused engine with at most 256 node types (engine.rs:866-868), while ShardHandle ranges and sampled edge IDs remain u64; no reachable-ID fallback exists. The added equivalence tests use only IDs below 41. No PR conversation comment containing the required change-map marker was available, so its design rationale and verification claims cannot be confirmed. This is still a changed-code behavior-preservation defect.
Resolution
Use a collision-free key that retains the full u64 ID and type, or disable per-seed recording when any reachable ID is at least 2^56. Document the compatibility condition in the change map and add a release-compatible regression test with high valid IDs that compares per-seed and legacy output.
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Bound retained sampling buffers in SamplingState::reset.
PooledState::drop resets and returns each SamplingState to StatePool, but reset does not clear or shrink raw_edges or commit_scratch. A large layer can therefore retain high-water capacities across requests.
order, res, and new_edges scale with the layer edge count. bounds, cursor, and new_bounds scale with seed slots and can also retain large capacities. Apply the existing RESET_THRESHOLD policy to these vectors.
♻️ Suggested fix
for m in &mut self.seed_maps {
m.clear();
if m.capacity() > SEED_MAP_KEEP {
m.shrink_to(SEED_MAP_KEEP);
}
}
+ for buf in &mut self.raw_edges {+ buf.clear();+ if buf.capacity() > RESET_THRESHOLD {+ buf.shrink_to(RESET_THRESHOLD);+ }+ }+ let cs = &mut self.commit_scratch;+ for v in [&mut cs.res, &mut cs.bounds, &mut cs.cursor, &mut cs.new_bounds] {+ v.clear();+ if v.capacity() > RESET_THRESHOLD {+ v.shrink_to(RESET_THRESHOLD);+ }+ }+ cs.order.clear();+ if cs.order.capacity() > RESET_THRESHOLD {+ cs.order.shrink_to(RESET_THRESHOLD);+ }+ cs.new_edges.clear();+ if cs.new_edges.capacity() > RESET_THRESHOLD {+ cs.new_edges.shrink_to(RESET_THRESHOLD);+ }
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 4 remain after this review.
DiskGraphEngine::legacy_recording, sample_with_cutoff, process_layer, new commit_edges, emit_edge, record_edge, node_key
Standalone layers commit edges through per-seed maps: bucket by seed, dedup per seed, assign ids and emit in edge order, write ids back. Seeds are registered in per-seed maps; commit_edges is sized by the request's seed slots; map-phase edge buffers are reused; record_edge shares emit_edge; the flag keeps the previous path. A request errors once one type reaches 2^31 local ids.
The per-edge global index_map lookup was ~70% of deep-hop sampling CPU; output stays identical. The 2^31 check turns a silent id collision with the pending marker into an error.
Per-seed maps, slot bases, commit scratch and edge buffers live in the pooled state; RawEdge moved here. Reset caps edge buffers and scratch at RESET_THRESHOLD, keeps seed maps only while their total capacity fits that budget, and keeps surplus edge buffers when a type has fewer shards. The index_map doc says it holds only seeds under per-seed recording.
Reuse across requests and layers without one heavy request pinning its peak memory.
The test shard builder takes arbitrary per-node adjacency. Equivalence vs the previous path is checked with one node type, and with two node types (A→B, B→A) and seeds of both types at overlapping indices, at 3/40/700 seeds and 1 and 4 threads.
Proves identical ids, edges, nodes and stats, including the per-seed-type slot offsets and the type bits in node_key.
Use one shared retention budget for per-seed maps.
A large request can grow up to 1024 seed maps, each retaining at least 16,384 element slots after reset. That is over 16 million retained slots per SamplingState, before hash-table metadata. StatePool has no size limit, so concurrent states can retain this memory until the pool is dropped.
Keep surplus edge buffers when the shard count shrinks.
resize_with(shards.len(), Vec::new) truncates buffers when the previous node type had more shards than the current one. The zip also drops the excess buffers. The next type with more shards then allocates new empty buffers. A heterogeneous graph with types of different shard counts loses the buffer reuse that this PR adds. The loss repeats for every type and layer.
Split off the surplus buffers before the map phase. Append them back to state.raw_edges after each reassignment of state.raw_edges at Lines 2776 and 2797. Do not pass the surplus buffers to commit_edges or to the reduce loop, because they can hold stale edges.
♻️ Proposed change
let mut buffers = std::mem::take(&mut state.raw_edges);
+ let spare = if buffers.len() > shards.len() {+ buffers.split_off(shards.len())+ } else {+ Vec::new()+ };
buffers.resize_with(shards.len(), Vec::new);
Then, in both places that assign state.raw_edges = raw_edges_results;, add this line directly after the assignment:
Reviewing files that changed from the base of the PR and between 954c87a and 16465c8.
📒 Files selected for processing: 2
packages/diskgraph/diskgraph-engine/src/engine.rs
packages/diskgraph/diskgraph-engine/src/state.rs
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
Base automatically changed from
vedaant/dg-sampler-type-order to
mainOctober 27, 2026 21:06
Each standalone sampling layer committed its edges one at a time through
one hash map keyed by (seed, node) over every node sampled so far, which
grows to millions of entries and misses cache on nearly every lookup.
Edges are now grouped by seed and deduplicated in a small per-seed map,
then local ids are assigned and edges emitted in the original order, so
ids, edge order and stats are unchanged. Per-shard edge buffers are
reused across layers and requests and, like index_map, capped on reset.
`DiskGraphEngine::legacy_recording`
selects the previous path.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Approach to split by seed looks solid. The id and edge-order logic appears correct. In addition to the inline comments, small suggestion: Add a benchmark on the training path. The numbers here use pooled DiskGraphEngine::sample. Training goes through runtime_batch_iterator → sample_runtime_batch, which creates a fresh SamplingState::default() for every batch. The trainer's auto-config (diskgraph_nested_trainer.py:919) selects standalone mode for larger subgraphs, which is the case this PR targets. On that path, reusing pooled buffers never applies, every batch grows about b seed maps from empty, and LP requests add several seed types. Could you add one measurement for the stack through sample_runtime_batch, ideally with an LP plan?
2882-2882: Minor · Performance & Scalability · Quick win
Use the active slot count for commit loops.
SamplingState::reset retains up to SEED_MAPS_KEEP maps. A later request can use fewer slots, but commit_edges uses state.seed_maps.len() and iterates every retained map for each layer and source type. Store the current request’s slot count when sample_with_cutoff computes it, and use that count for the bounds, cursor, and new_bounds loops. Keep seed_maps.len() as the capacity.
Buffer reuse loses buffers when the shard count decreases.
buffers.resize_with only grows the vector. Assume state.raw_edges holds more buffers than shards.len(). The zip with shard_groups then stops at the shorter side. buffers.into_par_iter() drops the unused buffers together with their capacity. Line 2776 and Line 2797 then store the shorter raw_edges_results back in state.raw_edges.
work_queue is processed by type ID, so a type with few shards can follow a type with many shards. This happens in every layer. The next type with many shards then allocates empty buffers again. The reuse of per-shard buffers then fails for graphs where node types have different shard counts.
Keep the surplus buffers and return them after the commit. Add the same line at Line 2776 (before continue) and at Line 2797.
♻️ Proposed fix
let mut buffers = std::mem::take(&mut state.raw_edges);
- buffers.resize_with(shards.len(), Vec::new);+ let spare = if buffers.len() > shards.len() {+ buffers.split_off(shards.len())+ } else {+ buffers.resize_with(shards.len(), Vec::new);+ Vec::new()+ };
// After each `state.raw_edges = raw_edges_results;` (Line 2776 and Line 2797):
state.raw_edges.extend(spare);
No existing legacy comparison uses multiple seed types. repeated_runs_give_identical_output exercises overlapping seed index 0 for lhs and rhs, including the lhs-to-rhs relation, but it checks only repeated per-seed consistency.
Compare the non-fused run with legacy recording. Capture every type state, not only rhs: MockOutput::edges does not include node types, so the current proposed assertion could miss a type-key assignment that changes the lhs state while preserving numeric edge IDs.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
A cleared map shrunk to a per-map cap still keeps its rounded-up bucket
array, so 1024 retained maps could pin ~0.5 GiB per pooled state. Keep
maps only while their summed capacity fits RESET_THRESHOLD and free the
rest.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A pooled state keeps up to 1024 cleared seed maps from earlier requests;
the bucket and prefix-sum loops now cover only the current request's slots.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Node types with different shard counts no longer drop the extra map-phase
buffers, so a later type with more shards reuses them instead of
reallocating.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
commit_edges tags pending entries with bit 31, so a larger local id would
be misread as an edge index. Fail the request instead.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A->B and B->A relations with seeds of both types at overlapping indices,
3/40/700 seeds, 1 and 4 threads: covers seed_slot_base and the type bits
of node_key against the legacy path.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Formats only lines this change adds; the surrounding map-phase block was
already unformatted and is left as is.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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
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.
GitHub Actions
Inactive
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Summary
This PR makes multi-hop DiskGraph sampling 1.3–2× faster per thread by changing how each layer records sampled edges. The output is identical to the current path.
Problem. Deep sampling plans spend most of their time after the edges are read:
Cause. Each layer runs in two phases. The map phase reads adjacency in parallel, one shard per task. The commit phase then runs serially: for every sampled edge,
record_edgelooks up(seed, node)in oneindex_mapper node type, which holds every node sampled so far for every seed in the request. At b=512 and 5 hops that map reaches millions of entries, so nearly every lookup misses cache. Profiling 5-hop sampling on one thread: commit (record_edgeand its hash-map lookups, inserts, rehashes and page faults) took about 70% of the CPU and the parallel map phase about 16%.Change.
commit_edgesreplaces the per-edge loop for the standalone (non-fused) path:seed_maps[slot], about 10k nodes at 6 hops, which stays in cache). Each edge resolves to an existing local id, a new node, or a repeat of an earlier edge in this layer.record_edgeproduced.The per-shard map-phase edge buffers are also kept in
SamplingStateand reused across node types, layers and requests instead of being reallocated.DiskGraphEngine::legacy_recording = trueselects the previous path; the fused path is unchanged.Back-compat
fused=false)legacy_recording = trueRESET_THRESHOLD(262k entries, about 4 MiB); the rest are freed. Edge buffers and commit scratch are cleared and shrunk toRESET_THRESHOLD, asindex_mapalready is. One heavy request does not pin its peak memory in the pool.TypeState::index_mapMeasured results
Indexes: synthetic relational graphs with 7 node types (2 entity tables and 5 event tables), 106.5M rows (5.3 GB index, 5 shards per type) and 1.01B rows (50 GB, 16 shards per type). Workload: temporal most-recent sampling with fanouts 24, 8, 6, 8, 6, 8 (first h used), 5,000 seeds sharing one anchor time, b=512 per request, warm page cache. Measured with a kumo-stress-tests ctxbench-style driver that calls
DiskGraphEngine::samplein-process (andfetch_features_for_typefor the feature-fetch columns). Each cell is the better of two interleaved rounds. Machine: 64 vCPU, 249 GiB RAM, shared with other jobs, so expect about ±10%.Sampling seeds/s,
main→ this PR:50B results for the full stack are in #102.
Training path (
sample_runtime_batch, link prediction)The table above uses pooled
DiskGraphEngine::sample. Training callssample_runtime_batch, which builds a freshSamplingStateper batch, so no maps or buffers carry over between batches. This run goes through that call on the 1B index with a two-tower link-prediction stage table (200k rows of customer → merchant purchases, entity = customers, target = merchants, one engine negative per row), so each batch seeds three groups: 512 customers, 512 positive merchants and 512 negative merchants. Same fanouts, temporal most-recent; features skipped through the feature projection unless noted. Per-batch time, better of three interleaved rounds of 4 batches, same machine (±10%).legacy_recording = trueis the previous path in the same binary; both paths returned identical node counts in every cell.At 2 hops this PR alone is 10–14% slower at 1 thread (31,250 vs 36,496 on 100M; 27,473 vs 30,675 on 1B): the commit phase is small there, so the extra bucketing pass is not paid back. #102 recovers it (0.96–0.99× of
main). Peak RSS rises by 0.0–0.3 GB.Known limitations
Testing Done
per_seed_recording_matches_legacy_recording: a random 41-node graph with repeated neighbours, temporal and uniform modes, both edge directions, repeated seeds, 1–700 seeds per request at 4 threads, and pooled state reused across requests; edges, nodes and stats equal the previous path.per_seed_recording_matches_legacy_recording_across_seed_types: two node types (A→B, B→A) with seeds of both types at overlapping indices, 3/40/700 seeds, 1 and 4 threads; covers the per-type slot offsets and the type bits of the per-seed key.tests/per_seed_recording_real_graph.rs(#[ignore], runs withDISKGRAPH_VERIFY_INDEX=<index>): 6 batches of 512 seeds, alternating a shared and per-seed anchors, 2–6 hops; checks this path equals the previous one and that the previous one is identical across fresh states. Passed on both indexes above (up to 14.4M edges per hop count).reset_caps_edge_buffers_and_commit_scratch: after a request that grew the edge buffers, commit scratch and per-seed maps pastRESET_THRESHOLD, the next acquired state holds at most that capacity in each.cargo test -p diskgraph-engine,-p diskgraph-ingest --test seed_table_ingest --test resolution_runtime_batchand-p diskgraph-pybindpass;cargo clippy --workspace --all-targets -- -D warningsis clean.Checklist
🤖 Generated with Claude Code