From c9d09d2acf3e5fc48794ca42be074d285681e28b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 07:48:10 +0000 Subject: [PATCH 1/9] diskgraph: record sampled edges through per-seed maps 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 --- .../diskgraph/diskgraph-engine/src/engine.rs | 195 +++++++++++++++--- .../diskgraph/diskgraph-engine/src/state.rs | 82 ++++++++ .../diskgraph-engine/src/tests_unit.rs | 139 ++++++++++--- .../tests/per_seed_recording_real_graph.rs | 121 +++++++++++ 4 files changed, 477 insertions(+), 60 deletions(-) create mode 100644 packages/diskgraph/diskgraph-engine/tests/per_seed_recording_real_graph.rs diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 6d91107fa..4d8399fc7 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -5,7 +5,7 @@ use crate::edge_unpacker::EdgeUnpacker; use crate::errors::{EngineError, Result}; use crate::output::OutputHandler; -use crate::state::{LookupKey, NodeRecord, ResolvedProjection, SamplingState}; +use crate::state::{LookupKey, NodeRecord, RawEdge, ResolvedProjection, SamplingState}; use diskgraph_core::sampling::*; use diskgraph_core::sortkey::SortKeyDecoder; use diskgraph_core::tt_builder::{Value, WEIGHT_COL}; @@ -462,6 +462,17 @@ pub struct DiskGraphEngine { /// Uses a per-instance pool (not the global pool) to avoid /// deadlocks when PyTorch DataLoader forks worker processes. pub thread_pool: rayon::ThreadPool, + /// Record edges through the global `index_map`; same output. + pub legacy_recording: bool, +} + +const NEW_NODE: u32 = u32::MAX; +const PENDING: u32 = 1 << 31; + +#[inline] +fn node_key(type_id: usize, global_id: u64) -> u64 { + debug_assert!(global_id < 1 << 56); + (global_id << 8) | type_id as u64 } struct CompiledSamplingPlan { @@ -473,18 +484,6 @@ struct CompiledLayerSpec { default_fanout: Option, } -/// A discovered edge before it is formally committed to the SamplingState. -struct RawEdge { - src_local_id: u32, - dst_global_id: u64, - dst_type_id: usize, - rel_id: u16, - seed_index: u32, - /// Root seed type id, propagated from the source node so the destination - /// node can resolve its timestamp anchor against the correct per-type array. - seed_type_id: u32, - swap_edge_direction: bool, -} impl DiskGraphEngine { /// Build an engine directly from an already-opened `GraphIndexReader`. @@ -860,9 +859,14 @@ impl DiskGraphEngine { local_to_merged_rel, random_seed, thread_pool, + legacy_recording: false, } } + fn per_seed_recording(&self, fused: bool) -> bool { + !fused && !self.legacy_recording && self.readers.len() <= 256 + } + pub fn get_type_name_mapping(&self) -> TypeNameMapping { let edge_relations = self .id_to_rel_name @@ -2120,6 +2124,21 @@ impl DiskGraphEngine { let fused = request.flags.fused_feature_fetching; + let per_seed = self.per_seed_recording(fused); + if per_seed { + state.seed_slot_base.resize(num_types, 0); + let mut slots = 0u32; + for (seed_type, seed_ids) in &request.seed_nodes { + if let Some(&t) = self.node_type_to_id.get(seed_type) { + state.seed_slot_base[t] = slots; + slots += seed_ids.len() as u32; + } + } + if state.seed_maps.len() < slots as usize { + state.seed_maps.resize_with(slots as usize, Default::default); + } + } + // 2. Initialize Seed Layer (Hop 0) // This ensures seed nodes get Local IDs 0..K and are preserved in order. for (seed_type, seed_ids) in &request.seed_nodes { @@ -2139,6 +2158,10 @@ impl DiskGraphEngine { if !seed_state.index_map.contains_key(&key) { let local_id = seed_state.nodes.len() as u32; seed_state.index_map.insert(key, local_id); + if per_seed { + let slot = (state.seed_slot_base[seed_type_id] + seed_index) as usize; + state.seed_maps[slot].insert(node_key(seed_type_id, global_id), local_id); + } // Cache Node (Hop 0). A seed is its own root, so its // root seed type is its own type id. @@ -2881,11 +2904,17 @@ impl DiskGraphEngine { // Map Phase: Parallel Shard Traversal // Each thread processes one shard to emit all sampled RawEdge + let mut buffers = std::mem::take(&mut state.raw_edges); + buffers.resize_with(shards.len(), Vec::new); let raw_edges_results: Vec> = self.thread_pool.install(|| shard_groups .into_par_iter() + .zip(buffers.into_par_iter()) .enumerate() - .filter(|(_, nodes)| !nodes.is_empty()) - .map(|(shard_idx, mut shard_nodes)| { + .map(|(shard_idx, (mut shard_nodes, mut local_raw_edges))| { + local_raw_edges.clear(); + if shard_nodes.is_empty() { + return local_raw_edges; + } let shard = &shards[shard_idx]; // STABLE SORT: Sort by global_id, then seed_index to ensure @@ -2894,7 +2923,6 @@ impl DiskGraphEngine { a.global_id.cmp(&b.global_id).then(a.seed_index.cmp(&b.seed_index)) }); - let mut local_raw_edges = Vec::new(); let mut current_segments = Vec::new(); let mut last_id = None; @@ -2968,8 +2996,14 @@ impl DiskGraphEngine { .collect() ); + if self.per_seed_recording(fused) { + self.commit_edges(layer_idx, &raw_edges_results, state, output)?; + state.raw_edges = raw_edges_results; + continue; + } + // Reduce Phase: Sequential Commitment - for raw_edges in raw_edges_results { + for raw_edges in &raw_edges_results { for edge in raw_edges { self.record_edge( layer_idx, @@ -2985,6 +3019,7 @@ impl DiskGraphEngine { )?; } } + state.raw_edges = raw_edges_results; } Ok(()) @@ -3054,6 +3089,118 @@ impl DiskGraphEngine { Ok(()) } + /// Same result as `record_edge` per edge, deduplicated per seed. + fn commit_edges( + &self, + layer_idx: usize, + chunks: &[Vec], + state: &mut SamplingState, + output: &mut impl OutputHandler, + ) -> Result<()> { + let n: usize = chunks.iter().map(Vec::len).sum(); + if n >= PENDING as usize - 1 { + return Err(EngineError::Internal(format!("too many edges in one layer: {n}"))); + } + let mut cs = std::mem::take(&mut state.commit_scratch); + let base = state.seed_slot_base.clone(); + let slot = |e: &RawEdge| (base[e.seed_type_id as usize] + e.seed_index) as usize; + let slots = state.seed_maps.len(); + + cs.bounds.clear(); + cs.bounds.resize(slots + 1, 0); + for e in chunks.iter().flatten() { + cs.bounds[slot(e) + 1] += 1; + } + for s in 0..slots { + cs.bounds[s + 1] += cs.bounds[s]; + } + cs.cursor.clear(); + cs.cursor.extend_from_slice(&cs.bounds[..slots]); + cs.order.resize(n, (0, 0)); + for (i, e) in chunks.iter().flatten().enumerate() { + let cur = &mut cs.cursor[slot(e)]; + cs.order[*cur as usize] = (node_key(e.dst_type_id, e.dst_global_id), i as u32); + *cur += 1; + } + + cs.res.resize(n, 0); + cs.new_edges.clear(); + cs.new_bounds.clear(); + for s in 0..slots { + let map = &mut state.seed_maps[s]; + for &(key, i) in &cs.order[cs.bounds[s] as usize..cs.bounds[s + 1] as usize] { + match map.entry(key) { + std::collections::hash_map::Entry::Occupied(o) => cs.res[i as usize] = *o.get(), + std::collections::hash_map::Entry::Vacant(v) => { + v.insert(PENDING | i); + cs.res[i as usize] = NEW_NODE; + cs.new_edges.push((key, i)); + } + } + } + cs.new_bounds.push(cs.new_edges.len() as u32); + } + + for (i, e) in chunks.iter().flatten().enumerate() { + let r = cs.res[i]; + let dst_local_id = if r == NEW_NODE { + let ts = state.type_states.entry(e.dst_type_id).or_default(); + let id = ts.nodes.len() as u32; + let node = NodeRecord { + global_id: e.dst_global_id, + seed_index: e.seed_index, + local_id: id, + seed_type_id: e.seed_type_id, + }; + ts.nodes.push(node); + ts.next_frontier.push(node); + state.stats.record_node(layer_idx + 1, e.dst_type_id); + cs.res[i] = id; + id + } else if r & PENDING != 0 { + cs.res[(r & !PENDING) as usize] + } else { + r + }; + self.emit_edge(layer_idx, e.src_local_id, dst_local_id, e.rel_id, e.swap_edge_direction, state, output)?; + } + + let mut lo = 0; + for (s, &hi) in cs.new_bounds.iter().enumerate() { + let map = &mut state.seed_maps[s]; + for &(key, i) in &cs.new_edges[lo..hi as usize] { + map.insert(key, cs.res[i as usize]); + } + lo = hi as usize; + } + state.commit_scratch = cs; + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + fn emit_edge( + &self, + layer_idx: usize, + src_local_id: u32, + dst_local_id: u32, + rel_id: u16, + swap_edge_direction: bool, + state: &mut SamplingState, + output: &mut impl OutputHandler, + ) -> Result<()> { + let (actual_rel, actual_src, actual_dst) = if swap_edge_direction { + match self.reverse_rel_map.get(rel_id as usize) { + Some(Some(rev_id)) => (*rev_id as u16, dst_local_id, src_local_id), + _ => (rel_id, src_local_id, dst_local_id), + } + } else { + (rel_id, src_local_id, dst_local_id) + }; + output.handle_edge(layer_idx, actual_rel, actual_src, actual_dst)?; + state.stats.record_edge(layer_idx, actual_rel as usize); + Ok(()) + } + #[allow(clippy::too_many_arguments)] fn record_edge( &self, @@ -3099,19 +3246,7 @@ impl DiskGraphEngine { }; // Apply edge direction reversal if requested (PyG CSC convention). - let (actual_rel, actual_src, actual_dst) = if swap_edge_direction { - match self.reverse_rel_map.get(rel_id as usize) { - Some(Some(rev_id)) => (*rev_id as u16, dst_local_id, src_local_id), - _ => (rel_id, src_local_id, dst_local_id), - } - } else { - (rel_id, src_local_id, dst_local_id) - }; - - output.handle_edge(layer_idx, actual_rel, actual_src, actual_dst)?; - state.stats.record_edge(layer_idx, actual_rel as usize); - - Ok(()) + self.emit_edge(layer_idx, src_local_id, dst_local_id, rel_id, swap_edge_direction, state, output) } } diff --git a/packages/diskgraph/diskgraph-engine/src/state.rs b/packages/diskgraph/diskgraph-engine/src/state.rs index d5e4bba36..38bc4c62f 100644 --- a/packages/diskgraph/diskgraph-engine/src/state.rs +++ b/packages/diskgraph/diskgraph-engine/src/state.rs @@ -250,8 +250,47 @@ pub struct SamplingState { /// Statistics for the current run. pub stats: SamplingStats, + + /// slot = seed_slot_base[seed_type_id] + seed_index + pub seed_slot_base: Vec, + pub seed_maps: Vec>, + pub commit_scratch: CommitScratch, + pub raw_edges: Vec>, +} + +/// A discovered edge before it is formally committed to the SamplingState. +pub struct RawEdge { + pub src_local_id: u32, + pub dst_global_id: u64, + pub dst_type_id: usize, + pub rel_id: u16, + pub seed_index: u32, + /// Root seed type id, propagated from the source node so the destination + /// node can resolve its timestamp anchor against the correct per-type array. + pub seed_type_id: u32, + pub swap_edge_direction: bool, +} + +#[derive(Default)] +pub struct CommitScratch { + pub bounds: Vec, + pub cursor: Vec, + pub order: Vec<(u64, u32)>, + pub res: Vec, + pub new_edges: Vec<(u64, u32)>, + pub new_bounds: Vec, +} + +fn clear_and_cap(v: &mut Vec, cap: usize) { + v.clear(); + if v.capacity() > cap { + v.shrink_to(cap); + } } +const SEED_MAP_KEEP: usize = 1 << 14; +const SEED_MAPS_KEEP: usize = 1 << 10; + const RESET_THRESHOLD: usize = 1 << 18; // ~262k items impl Default for SamplingState { @@ -269,6 +308,10 @@ impl SamplingState { reservoir_scratch: Vec::with_capacity(128), decode_scratch: diskgraph_rowpack::RowPackDecodeScratch::default(), stats: SamplingStats::new(), + seed_slot_base: Vec::new(), + seed_maps: Vec::new(), + commit_scratch: CommitScratch::default(), + raw_edges: Vec::new(), } } @@ -281,6 +324,25 @@ impl SamplingState { self.key_scratch.clear(); self.reservoir_scratch.clear(); self.stats.reset(); + self.seed_slot_base.clear(); + let per_buffer = RESET_THRESHOLD / self.raw_edges.len().max(1); + for v in &mut self.raw_edges { + clear_and_cap(v, per_buffer); + } + let cs = &mut self.commit_scratch; + clear_and_cap(&mut cs.bounds, RESET_THRESHOLD); + clear_and_cap(&mut cs.cursor, RESET_THRESHOLD); + clear_and_cap(&mut cs.order, RESET_THRESHOLD); + clear_and_cap(&mut cs.res, RESET_THRESHOLD); + clear_and_cap(&mut cs.new_edges, RESET_THRESHOLD); + clear_and_cap(&mut cs.new_bounds, RESET_THRESHOLD); + self.seed_maps.truncate(SEED_MAPS_KEEP); + for m in &mut self.seed_maps { + m.clear(); + if m.capacity() > SEED_MAP_KEEP { + m.shrink_to(SEED_MAP_KEEP); + } + } // Scratch buffer is inherently reusable } /// Access or create the state for a specific node type. @@ -357,6 +419,26 @@ impl Drop for PooledState { mod tests { use super::*; + #[test] + fn reset_caps_edge_buffers_and_commit_scratch() { + let pool = StatePool::new(); + { + let mut state = pool.acquire(); + state.raw_edges = (0..4).map(|_| Vec::with_capacity(RESET_THRESHOLD * 2)).collect(); + let cs = &mut state.commit_scratch; + cs.order.reserve(RESET_THRESHOLD * 4); + cs.res.reserve(RESET_THRESHOLD * 4); + cs.new_edges.reserve(RESET_THRESHOLD * 4); + } + let state = pool.acquire(); + let edge_capacity: usize = state.raw_edges.iter().map(Vec::capacity).sum(); + assert!(edge_capacity <= RESET_THRESHOLD); + let cs = &state.commit_scratch; + assert!(cs.order.capacity() <= RESET_THRESHOLD); + assert!(cs.res.capacity() <= RESET_THRESHOLD); + assert!(cs.new_edges.capacity() <= RESET_THRESHOLD); + } + #[test] fn test_pool_recycling() { let pool = StatePool::new(); diff --git a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs index ffb944a2a..9bf6e72b3 100644 --- a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs +++ b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs @@ -1373,7 +1373,15 @@ fn build_sort_key_shard( max_node: u8, key_bits: u8, ) -> (ShardReader, RelationLayout, diskgraph_rowpack::RowPackContext) { - // Node feature context (one trivial column), mirroring TestGraph. + let (ctx, node_data) = trivial_node_rowpack(); + let mut adjacency = vec![Vec::new(); max_node as usize + 1]; + adjacency[0] = entries.to_vec(); + let (reader, layout) = write_adjacency_shard(dir, &adjacency, key_bits, &node_data); + (reader, layout, ctx) +} + +/// Node feature context (one trivial column), mirroring TestGraph. +fn trivial_node_rowpack() -> (diskgraph_rowpack::RowPackContext, Vec) { let cols = vec![RowPackColumn { name: "feat_0".into(), logical_type: IngestedType::Int64, @@ -1391,7 +1399,17 @@ fn build_sort_key_shard( encoder .encode_row_into(&row_val, &mut node_data, &mut scratch) .unwrap(); + (ctx, node_data) +} +/// Writes one shard whose node `i` has the `(dst, sort_key)` entries in +/// `adjacency[i]` (rel 0, 8-bit dst ids, at most 63 nodes). +fn write_adjacency_shard( + dir: &std::path::Path, + adjacency: &[Vec<(u8, u64)>], + key_bits: u8, + node_data: &[u8], +) -> (ShardReader, RelationLayout) { let entry_bytes = 1 + key_bits as usize / 8; let layout = RelationLayout { relation_id: 0, @@ -1402,52 +1420,42 @@ fn build_sort_key_shard( dest_sort_key_mask: Some(u64::MAX >> (64 - key_bits)), }; - // edges.bin: node 0 has `entries.len()` neighbors; nodes 1..=max_node have - // zero edges. Each node block: [count u8][len u16][entry_bytes*count edge - // bytes][node_data]. + // Each node block: [count u8][len u16][entry_bytes*count edge bytes][node_data]. let mut f = File::create(dir.join(EDGES_FILENAME)).unwrap(); - let node0_edge_bytes: Vec = entries - .iter() - .flat_map(|&(dst, sort_key)| { - let entry = (dst as u128) | ((sort_key as u128) << 8); - entry.to_le_bytes()[..entry_bytes].to_vec() - }) - .collect(); - f.write_all(&[entries.len() as u8]).unwrap(); - f.write_all(&(node_data.len() as u16).to_le_bytes()).unwrap(); - f.write_all(&node0_edge_bytes).unwrap(); - f.write_all(&node_data).unwrap(); - let node0_len = (1 + 2 + node0_edge_bytes.len() + node_data.len()) as u64; - - let empty_node_len = (1 + 2 + node_data.len()) as u64; - for _ in 1..=max_node { - f.write_all(&[0x00]).unwrap(); + let mut lens = Vec::new(); + for entries in adjacency { + let edge_bytes: Vec = entries + .iter() + .flat_map(|&(dst, sort_key)| { + let entry = (dst as u128) | ((sort_key as u128) << 8); + entry.to_le_bytes()[..entry_bytes].to_vec() + }) + .collect(); + f.write_all(&[entries.len() as u8]).unwrap(); f.write_all(&(node_data.len() as u16).to_le_bytes()).unwrap(); - f.write_all(&node_data).unwrap(); + f.write_all(&edge_bytes).unwrap(); + f.write_all(node_data).unwrap(); + lens.push((1 + 2 + edge_bytes.len() + node_data.len()) as u64); } f.flush().unwrap(); let edges_len = f.metadata().unwrap().len(); - // offsets.bin: (max_node + 1) nodes => (max_node + 2) offsets. - let num_nodes = max_node as u64 + 1; + let num_nodes = adjacency.len() as u64; let header = OffsetFileHeader::new(num_nodes, num_nodes + 1); let mut offsets = bincode::serialize(&header).unwrap(); offsets.extend_from_slice(&((num_nodes + 1) as u16).to_le_bytes()); offsets.extend_from_slice(&0u64.to_le_bytes()); // base offset - offsets.extend_from_slice(&(node0_len as u16).to_le_bytes()); - let mut running = node0_len; - for _ in 1..=max_node { - offsets.extend_from_slice(&(empty_node_len as u16).to_le_bytes()); - running += empty_node_len; + for &len in &lens { + offsets.extend_from_slice(&(len as u16).to_le_bytes()); } - assert_eq!(running, edges_len, "offsets must sum to edge file size"); + assert_eq!(lens.iter().sum::(), edges_len, "offsets must sum to edge file size"); let mut fo = File::create(dir.join(OFFSETS_FILENAME)).unwrap(); fo.write_all(&offsets).unwrap(); fo.flush().unwrap(); let reader = ShardReader::open(dir, Arc::new(vec![layout.clone()]), true).unwrap(); - (reader, layout, ctx) + (reader, layout) } /// Run a `TemporalMostRecent` sample over a two-seed-type graph and return the @@ -2497,6 +2505,77 @@ fn test_finalize_preserves_anchor_major_seed_prefix() { ); } +type SampleSnapshot = (Vec<(usize, u16, u32, u32)>, Vec>, Vec>, Vec>); + +#[test] +fn per_seed_recording_matches_legacy_recording() { + use rand::{Rng, SeedableRng}; + let dir = TempDir::new().unwrap(); + let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(11); + let sk = SortKeyEncoder::new(&temporal_sortkey_meta()); + let n = 41u8; + let adjacency: Vec> = (0..n) + .map(|_| { + let mut e: Vec<(u8, u64)> = (0..rng.gen_range(0..30)) + .map(|_| (rng.gen_range(0..n), sk.encode_lower(rng.gen_range(0..1_000)))) + .collect(); + e.sort_by_key(|&(_, k)| std::cmp::Reverse(k)); + e + }) + .collect(); + let (ctx, node_data) = trivial_node_rowpack(); + let (reader, _) = write_adjacency_shard(dir.path(), &adjacency, 16, &node_data); + let reader = Arc::new(reader); + let mut sampler = DiskGraphEngine::new( + [("User".to_string(), vec![ShardHandle { reader, min_id: 0, max_id: n as u64 - 1 }])].into(), + [(0, ("Connects".to_string(), "User".to_string(), "User".to_string()))].into(), + [(0usize, SortKeyDecoder::new(&temporal_sortkey_meta()))].into(), + [("User".to_string(), RowPackDecoder::new(ctx).unwrap())].into(), + None, + diskgraph_core::time::TimeUnit::Seconds, + Some(42), + Some(2), + ); + + let requests: Vec = (0..6) + .map(|i| { + let mode = if i % 2 == 0 { SamplingMode::TemporalMostRecent } else { SamplingMode::Uniform }; + let seeds: Vec = (0..rng.gen_range(1..60)).map(|_| rng.gen_range(0..n as u64)).collect(); + let mut req = make_sampling_request_with_mode(seeds.clone(), vec![6, 4, 3, 5], mode); + req.seed_timestamp_unit = diskgraph_core::time::TimeUnit::Seconds; + req.seed_timestamps = Some([("User".to_string(), seeds.iter().map(|_| rng.gen_range(0..1_000)).collect())].into()); + req.flags.swap_edge_direction = i % 3 == 0; + req + }) + .collect(); + + let run = |sampler: &DiskGraphEngine| -> Vec { + let pool = crate::state::StatePool::new(); + requests + .iter() + .map(|req| { + let mut state = pool.acquire(); + let mut output = MockOutput { edges: vec![], feature_batches: vec![], schema: None, fused_batches: vec![] }; + sampler.sample(req, &mut state, &mut output).unwrap(); + let mut types: Vec<_> = state.type_states.iter().collect(); + types.sort_by_key(|(t, _)| **t); + let nodes = types + .iter() + .map(|(_, ts)| ts.nodes.iter().map(|r| (r.global_id, r.seed_index, r.local_id, r.seed_type_id)).collect()) + .collect(); + (output.edges, nodes, state.stats.nodes_per_type.clone(), state.stats.edges_per_type.clone()) + }) + .collect() + }; + + let per_seed = run(&sampler); + sampler.legacy_recording = true; + let legacy = run(&sampler); + assert_eq!(per_seed, legacy); + let (edges, nodes) = per_seed.iter().fold((0, 0), |(e, v), s| (e + s.0.len(), v + s.1[0].len())); + assert!(edges > nodes, "the graph must revisit nodes within a seed"); +} + #[test] fn repeated_runs_give_identical_output() { let (dir_lhs, dir_rhs) = (TempDir::new().unwrap(), TempDir::new().unwrap()); diff --git a/packages/diskgraph/diskgraph-engine/tests/per_seed_recording_real_graph.rs b/packages/diskgraph/diskgraph-engine/tests/per_seed_recording_real_graph.rs new file mode 100644 index 000000000..050732b56 --- /dev/null +++ b/packages/diskgraph/diskgraph-engine/tests/per_seed_recording_real_graph.rs @@ -0,0 +1,121 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +// All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! DISKGRAPH_VERIFY_INDEX= [DISKGRAPH_VERIFY_SEED_TYPE=customers] +//! [DISKGRAPH_VERIFY_ANCHOR=] cargo test --release -p diskgraph-engine +//! --test per_seed_recording_real_graph -- --ignored + +use diskgraph_core::sampling::*; +use diskgraph_engine::engine::DiskGraphEngine; +use diskgraph_engine::output::InMemoryOutput; +use diskgraph_engine::state::StatePool; +use std::collections::HashMap; + +const PLAN6: [u32; 6] = [24, 8, 6, 8, 6, 8]; + +fn request(seed_type: &str, seeds: &[u64], anchors: &[i64], hops: usize) -> BatchSamplingRequest { + let layers = PLAN6[..hops] + .iter() + .map(|&fanout| LayerSpec { + fanouts: HashMap::new(), + default_fanout: Some(FanoutSpec { + fanout, + mode: SamplingMode::TemporalMostRecent, + absent_key_policy: AbsentKeyPolicy::UniformFallback, + }), + }) + .collect(); + BatchSamplingRequest { + dataset_id: diskgraph_core::dataset::DatasetId::nil(), + seed_nodes: [(seed_type.to_string(), seeds.to_vec())].into(), + seed_timestamps: Some([(seed_type.to_string(), anchors.to_vec())].into()), + seed_timestamp_unit: diskgraph_core::time::TimeUnit::Seconds, + plan: SamplingPlan { layers }, + flags: SamplingFlags { + replace: false, + directed: true, + disjoint: true, + return_edge_id: false, + adaptive: false, + fused_feature_fetching: false, + swap_edge_direction: true, + }, + } +} + +#[test] +#[ignore = "needs DISKGRAPH_VERIFY_INDEX"] +fn per_seed_recording_matches_legacy_on_real_graph() { + let index = std::env::var("DISKGRAPH_VERIFY_INDEX").expect("DISKGRAPH_VERIFY_INDEX"); + let seed_type = + std::env::var("DISKGRAPH_VERIFY_SEED_TYPE").unwrap_or_else(|_| "customers".into()); + let anchor: i64 = + std::env::var("DISKGRAPH_VERIFY_ANCHOR").map_or(1_734_566_399, |s| s.parse().unwrap()); + let mut engine = DiskGraphEngine::load(std::path::Path::new(&index), Some(0), Some(4)).unwrap(); + let (lo, hi) = engine.node_type_id_range(&seed_type).unwrap(); + let (batch, batches) = (512usize, 6usize); + let n = batch * batches; + let seeds: Vec = (0..n as u64) + .map(|i| lo + i * (hi - lo) / n as u64) + .collect(); + let anchors: Vec = (0..n) + .map(|i| { + if (i / batch) % 2 == 0 { + anchor + } else { + anchor - (i as i64 * 7_919) % 31_536_000 + } + }) + .collect(); + + for hops in 2..=6 { + let mut runs = Vec::new(); + for legacy in [false, true, true] { + engine.legacy_recording = legacy; + let pool = StatePool::new(); + let mut out = Vec::new(); + for (s, a) in seeds.chunks(batch).zip(anchors.chunks(batch)) { + let mut state = pool.acquire(); + let mut output = InMemoryOutput::new(); + engine + .sample(&request(&seed_type, s, a, hops), &mut state, &mut output) + .unwrap(); + let mut types: Vec<_> = state.type_states.iter().collect(); + types.sort_by_key(|(t, _)| **t); + let nodes: Vec> = types + .iter() + .map(|(_, ts)| { + ts.nodes + .iter() + .map(|r| (r.global_id, r.seed_index, r.local_id, r.seed_type_id)) + .collect() + }) + .collect(); + let frontier: Vec> = types + .iter() + .map(|(_, ts)| ts.next_frontier.iter().map(|r| r.global_id).collect()) + .collect(); + out.push(( + output.edges, + nodes, + frontier, + state.stats.nodes_per_type.clone(), + state.stats.edges_per_type.clone(), + )); + } + runs.push(out); + } + let edges: usize = runs[0].iter().map(|r| r.0.len()).sum(); + assert!(edges > 0); + assert!( + runs[1] == runs[2], + "hops {hops}: legacy recording is not deterministic across states" + ); + assert!( + runs[0] == runs[1], + "hops {hops}: per-seed recording differs from legacy" + ); + println!("hops {hops}: {n} seeds, {edges} edges identical"); + } +} From 8d33df6e7ddff370f91041427b0339c5e9241a30 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:18:01 +0000 Subject: [PATCH 2/9] diskgraph: bound retained seed-map memory by one total budget 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 --- packages/diskgraph/diskgraph-engine/src/state.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/packages/diskgraph/diskgraph-engine/src/state.rs b/packages/diskgraph/diskgraph-engine/src/state.rs index 38bc4c62f..0b0250fc2 100644 --- a/packages/diskgraph/diskgraph-engine/src/state.rs +++ b/packages/diskgraph/diskgraph-engine/src/state.rs @@ -288,7 +288,6 @@ fn clear_and_cap(v: &mut Vec, cap: usize) { } } -const SEED_MAP_KEEP: usize = 1 << 14; const SEED_MAPS_KEEP: usize = 1 << 10; const RESET_THRESHOLD: usize = 1 << 18; // ~262k items @@ -337,11 +336,13 @@ impl SamplingState { clear_and_cap(&mut cs.new_edges, RESET_THRESHOLD); clear_and_cap(&mut cs.new_bounds, RESET_THRESHOLD); self.seed_maps.truncate(SEED_MAPS_KEEP); + let mut budget = RESET_THRESHOLD; for m in &mut self.seed_maps { m.clear(); - if m.capacity() > SEED_MAP_KEEP { - m.shrink_to(SEED_MAP_KEEP); + if m.capacity() > budget { + m.shrink_to(0); } + budget -= m.capacity(); } // Scratch buffer is inherently reusable } @@ -429,6 +430,7 @@ mod tests { cs.order.reserve(RESET_THRESHOLD * 4); cs.res.reserve(RESET_THRESHOLD * 4); cs.new_edges.reserve(RESET_THRESHOLD * 4); + state.seed_maps = (0..64).map(|_| GxHashMap::with_capacity(RESET_THRESHOLD / 16)).collect(); } let state = pool.acquire(); let edge_capacity: usize = state.raw_edges.iter().map(Vec::capacity).sum(); @@ -437,6 +439,9 @@ mod tests { assert!(cs.order.capacity() <= RESET_THRESHOLD); assert!(cs.res.capacity() <= RESET_THRESHOLD); assert!(cs.new_edges.capacity() <= RESET_THRESHOLD); + let map_capacity: usize = state.seed_maps.iter().map(GxHashMap::capacity).sum(); + assert!(map_capacity <= RESET_THRESHOLD); + assert!(map_capacity > 0); } #[test] From 12a45315a71856f066d36c8c9c9d603b7a719a98 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:18:10 +0000 Subject: [PATCH 3/9] diskgraph: size commit_edges by the request's seed slots 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 --- packages/diskgraph/diskgraph-engine/src/engine.rs | 7 ++++--- packages/diskgraph/diskgraph-engine/src/state.rs | 3 +++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 4d8399fc7..14ed065a2 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -2134,8 +2134,9 @@ impl DiskGraphEngine { slots += seed_ids.len() as u32; } } - if state.seed_maps.len() < slots as usize { - state.seed_maps.resize_with(slots as usize, Default::default); + state.seed_slots = slots as usize; + if state.seed_maps.len() < state.seed_slots { + state.seed_maps.resize_with(state.seed_slots, Default::default); } } @@ -3104,7 +3105,7 @@ impl DiskGraphEngine { let mut cs = std::mem::take(&mut state.commit_scratch); let base = state.seed_slot_base.clone(); let slot = |e: &RawEdge| (base[e.seed_type_id as usize] + e.seed_index) as usize; - let slots = state.seed_maps.len(); + let slots = state.seed_slots; cs.bounds.clear(); cs.bounds.resize(slots + 1, 0); diff --git a/packages/diskgraph/diskgraph-engine/src/state.rs b/packages/diskgraph/diskgraph-engine/src/state.rs index 0b0250fc2..4b92a65f1 100644 --- a/packages/diskgraph/diskgraph-engine/src/state.rs +++ b/packages/diskgraph/diskgraph-engine/src/state.rs @@ -253,6 +253,7 @@ pub struct SamplingState { /// slot = seed_slot_base[seed_type_id] + seed_index pub seed_slot_base: Vec, + pub seed_slots: usize, pub seed_maps: Vec>, pub commit_scratch: CommitScratch, pub raw_edges: Vec>, @@ -308,6 +309,7 @@ impl SamplingState { decode_scratch: diskgraph_rowpack::RowPackDecodeScratch::default(), stats: SamplingStats::new(), seed_slot_base: Vec::new(), + seed_slots: 0, seed_maps: Vec::new(), commit_scratch: CommitScratch::default(), raw_edges: Vec::new(), @@ -324,6 +326,7 @@ impl SamplingState { self.reservoir_scratch.clear(); self.stats.reset(); self.seed_slot_base.clear(); + self.seed_slots = 0; let per_buffer = RESET_THRESHOLD / self.raw_edges.len().max(1); for v in &mut self.raw_edges { clear_and_cap(v, per_buffer); From 3714bad92935838c47570f77cb1fc13a6e1d381d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:18:18 +0000 Subject: [PATCH 4/9] diskgraph: keep surplus edge buffers when a type has fewer shards 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 --- packages/diskgraph/diskgraph-engine/src/engine.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 14ed065a2..2c75fdefb 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -2906,6 +2906,7 @@ impl DiskGraphEngine { // Map Phase: Parallel Shard Traversal // Each thread processes one shard to emit all sampled RawEdge let mut buffers = std::mem::take(&mut state.raw_edges); + let spare = buffers.split_off(shards.len().min(buffers.len())); buffers.resize_with(shards.len(), Vec::new); let raw_edges_results: Vec> = self.thread_pool.install(|| shard_groups .into_par_iter() @@ -3000,6 +3001,7 @@ impl DiskGraphEngine { if self.per_seed_recording(fused) { self.commit_edges(layer_idx, &raw_edges_results, state, output)?; state.raw_edges = raw_edges_results; + state.raw_edges.extend(spare); continue; } @@ -3021,6 +3023,7 @@ impl DiskGraphEngine { } } state.raw_edges = raw_edges_results; + state.raw_edges.extend(spare); } Ok(()) From ffca9684e787ed9c58f06593797237961abfbaff Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:18:28 +0000 Subject: [PATCH 5/9] diskgraph: error when one type's local ids reach 2^31 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 --- packages/diskgraph/diskgraph-engine/src/engine.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 2c75fdefb..4d84438ec 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -3149,6 +3149,9 @@ impl DiskGraphEngine { let r = cs.res[i]; let dst_local_id = if r == NEW_NODE { let ts = state.type_states.entry(e.dst_type_id).or_default(); + if ts.nodes.len() >= PENDING as usize { + return Err(EngineError::Internal(format!("too many nodes of type {}", e.dst_type_id))); + } let id = ts.nodes.len() as u32; let node = NodeRecord { global_id: e.dst_global_id, From 1d5ac615de06066d94495fcb9bc45be69f0550cb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:18:28 +0000 Subject: [PATCH 6/9] diskgraph: note that index_map holds only seeds under per-seed recording Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-engine/src/state.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/diskgraph/diskgraph-engine/src/state.rs b/packages/diskgraph/diskgraph-engine/src/state.rs index 4b92a65f1..0ea0a261d 100644 --- a/packages/diskgraph/diskgraph-engine/src/state.rs +++ b/packages/diskgraph/diskgraph-engine/src/state.rs @@ -115,7 +115,8 @@ pub struct NodeRecord { /// State specific to a single Node Type (table). pub struct TypeState { - /// Maps (Seed, Global) -> LocalID + /// Maps (Seed, Global) -> LocalID. Holds only the seeds when edges are + /// recorded through `SamplingState::seed_maps`. pub index_map: GxHashMap, /// Ordered list of nodes discovered (Index = LocalID) From d113c7aef02f7a1f345cb22086424a5ec3117d98 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:20:00 +0000 Subject: [PATCH 7/9] diskgraph: test per-seed recording across two node and seed types 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 --- .../diskgraph-engine/src/tests_unit.rs | 130 +++++++++++++----- 1 file changed, 95 insertions(+), 35 deletions(-) diff --git a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs index 9bf6e72b3..969599e45 100644 --- a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs +++ b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs @@ -2507,14 +2507,10 @@ fn test_finalize_preserves_anchor_major_seed_prefix() { type SampleSnapshot = (Vec<(usize, u16, u32, u32)>, Vec>, Vec>, Vec>); -#[test] -fn per_seed_recording_matches_legacy_recording() { - use rand::{Rng, SeedableRng}; - let dir = TempDir::new().unwrap(); - let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(11); +fn random_adjacency(rng: &mut rand_chacha::ChaCha8Rng, n: u8) -> Vec> { + use rand::Rng; let sk = SortKeyEncoder::new(&temporal_sortkey_meta()); - let n = 41u8; - let adjacency: Vec> = (0..n) + (0..n) .map(|_| { let mut e: Vec<(u8, u64)> = (0..rng.gen_range(0..30)) .map(|_| (rng.gen_range(0..n), sk.encode_lower(rng.gen_range(0..1_000)))) @@ -2522,33 +2518,32 @@ fn per_seed_recording_matches_legacy_recording() { e.sort_by_key(|&(_, k)| std::cmp::Reverse(k)); e }) - .collect(); - let (ctx, node_data) = trivial_node_rowpack(); - let (reader, _) = write_adjacency_shard(dir.path(), &adjacency, 16, &node_data); - let reader = Arc::new(reader); - let mut sampler = DiskGraphEngine::new( - [("User".to_string(), vec![ShardHandle { reader, min_id: 0, max_id: n as u64 - 1 }])].into(), - [(0, ("Connects".to_string(), "User".to_string(), "User".to_string()))].into(), - [(0usize, SortKeyDecoder::new(&temporal_sortkey_meta()))].into(), - [("User".to_string(), RowPackDecoder::new(ctx).unwrap())].into(), - None, - diskgraph_core::time::TimeUnit::Seconds, - Some(42), - Some(2), - ); + .collect() +} - let requests: Vec = (0..6) - .map(|i| { - let mode = if i % 2 == 0 { SamplingMode::TemporalMostRecent } else { SamplingMode::Uniform }; - let seeds: Vec = (0..rng.gen_range(1..60)).map(|_| rng.gen_range(0..n as u64)).collect(); - let mut req = make_sampling_request_with_mode(seeds.clone(), vec![6, 4, 3, 5], mode); - req.seed_timestamp_unit = diskgraph_core::time::TimeUnit::Seconds; - req.seed_timestamps = Some([("User".to_string(), seeds.iter().map(|_| rng.gen_range(0..1_000)).collect())].into()); - req.flags.swap_edge_direction = i % 3 == 0; - req - }) - .collect(); +fn per_seed_request( + rng: &mut rand_chacha::ChaCha8Rng, + types: &[&str], + n: u8, + seeds: usize, + i: usize, +) -> BatchSamplingRequest { + use rand::Rng; + let mode = if i % 2 == 0 { SamplingMode::TemporalMostRecent } else { SamplingMode::Uniform }; + let mut req = make_sampling_request_with_mode(vec![], vec![6, 4, 3, 5], mode); + for layer in &mut req.plan.layers { + layer.default_fanout = layer.fanouts.drain().next().map(|(_, f)| f); + } + req.seed_nodes = types.iter().map(|t| (t.to_string(), (0..seeds).map(|_| rng.gen_range(0..n as u64)).collect())).collect(); + req.seed_timestamps = Some(types.iter().map(|t| (t.to_string(), (0..seeds).map(|_| rng.gen_range(0..1_000)).collect())).collect()); + req.seed_timestamp_unit = diskgraph_core::time::TimeUnit::Seconds; + req.flags.swap_edge_direction = i % 3 == 0; + req +} +/// Samples every request on both recording paths and asserts identical +/// edges, node records and stats. +fn assert_per_seed_matches_legacy(sampler: &mut DiskGraphEngine, requests: &[BatchSamplingRequest]) -> Vec { let run = |sampler: &DiskGraphEngine| -> Vec { let pool = crate::state::StatePool::new(); requests @@ -2567,15 +2562,80 @@ fn per_seed_recording_matches_legacy_recording() { }) .collect() }; - - let per_seed = run(&sampler); + sampler.legacy_recording = false; + let per_seed = run(sampler); sampler.legacy_recording = true; - let legacy = run(&sampler); + let legacy = run(sampler); assert_eq!(per_seed, legacy); + per_seed +} + +#[test] +fn per_seed_recording_matches_legacy_recording() { + use rand::SeedableRng; + let dir = TempDir::new().unwrap(); + let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(11); + let n = 41u8; + let adjacency = random_adjacency(&mut rng, n); + let (ctx, node_data) = trivial_node_rowpack(); + let (reader, _) = write_adjacency_shard(dir.path(), &adjacency, 16, &node_data); + let reader = Arc::new(reader); + let mut sampler = DiskGraphEngine::new( + [("User".to_string(), vec![ShardHandle { reader, min_id: 0, max_id: n as u64 - 1 }])].into(), + [(0, ("Connects".to_string(), "User".to_string(), "User".to_string()))].into(), + [(0usize, SortKeyDecoder::new(&temporal_sortkey_meta()))].into(), + [("User".to_string(), RowPackDecoder::new(ctx).unwrap())].into(), + None, + diskgraph_core::time::TimeUnit::Seconds, + Some(42), + Some(4), + ); + + let requests: Vec = + (0..6).map(|i| per_seed_request(&mut rng, &["User"], n, [1, 40, 700][i % 3], i)).collect(); + let per_seed = assert_per_seed_matches_legacy(&mut sampler, &requests); let (edges, nodes) = per_seed.iter().fold((0, 0), |(e, v), s| (e + s.0.len(), v + s.1[0].len())); assert!(edges > nodes, "the graph must revisit nodes within a seed"); } +#[test] +fn per_seed_recording_matches_legacy_recording_across_seed_types() { + use rand::SeedableRng; + let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(12); + let n = 41u8; + let (ctx, node_data) = trivial_node_rowpack(); + let dirs = [TempDir::new().unwrap(), TempDir::new().unwrap()]; + let readers: Vec> = dirs + .iter() + .map(|dir| Arc::new(write_adjacency_shard(dir.path(), &random_adjacency(&mut rng, n), 16, &node_data).0)) + .collect(); + let requests: Vec = + (0..6).map(|i| per_seed_request(&mut rng, &["A", "B"], n, [3, 40, 700][i % 3], i)).collect(); + for threads in [1, 4] { + let shard = |t: usize| vec![ShardHandle { reader: readers[t].clone(), min_id: 0, max_id: n as u64 - 1 }]; + let mut sampler = DiskGraphEngine::new( + [("A".to_string(), shard(0)), ("B".to_string(), shard(1))].into(), + [ + (0, ("AToB".to_string(), "A".to_string(), "B".to_string())), + (1, ("BToA".to_string(), "B".to_string(), "A".to_string())), + ] + .into(), + [(0usize, SortKeyDecoder::new(&temporal_sortkey_meta())), (1usize, SortKeyDecoder::new(&temporal_sortkey_meta()))].into(), + [ + ("A".to_string(), RowPackDecoder::new(ctx.clone()).unwrap()), + ("B".to_string(), RowPackDecoder::new(ctx.clone()).unwrap()), + ] + .into(), + None, + diskgraph_core::time::TimeUnit::Seconds, + Some(42), + Some(threads), + ); + let per_seed = assert_per_seed_matches_legacy(&mut sampler, &requests); + assert!(per_seed.iter().all(|s| s.1.len() == 2 && s.0.iter().any(|e| e.1 == 1))); + } +} + #[test] fn repeated_runs_give_identical_output() { let (dir_lhs, dir_rhs) = (TempDir::new().unwrap(), TempDir::new().unwrap()); From 9bac80ac72e429bd4296f2e205c749d0df60d41c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 21:26:17 +0000 Subject: [PATCH 8/9] diskgraph: rustfmt the per-seed recording lines 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 --- .../diskgraph/diskgraph-engine/src/engine.rs | 33 ++++- .../diskgraph/diskgraph-engine/src/state.rs | 8 +- .../diskgraph-engine/src/tests_unit.rs | 139 +++++++++++++++--- 3 files changed, 151 insertions(+), 29 deletions(-) diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 4d84438ec..f537e6f91 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -2136,7 +2136,9 @@ impl DiskGraphEngine { } state.seed_slots = slots as usize; if state.seed_maps.len() < state.seed_slots { - state.seed_maps.resize_with(state.seed_slots, Default::default); + state + .seed_maps + .resize_with(state.seed_slots, Default::default); } } @@ -3103,7 +3105,9 @@ impl DiskGraphEngine { ) -> Result<()> { let n: usize = chunks.iter().map(Vec::len).sum(); if n >= PENDING as usize - 1 { - return Err(EngineError::Internal(format!("too many edges in one layer: {n}"))); + return Err(EngineError::Internal(format!( + "too many edges in one layer: {n}" + ))); } let mut cs = std::mem::take(&mut state.commit_scratch); let base = state.seed_slot_base.clone(); @@ -3150,7 +3154,10 @@ impl DiskGraphEngine { let dst_local_id = if r == NEW_NODE { let ts = state.type_states.entry(e.dst_type_id).or_default(); if ts.nodes.len() >= PENDING as usize { - return Err(EngineError::Internal(format!("too many nodes of type {}", e.dst_type_id))); + return Err(EngineError::Internal(format!( + "too many nodes of type {}", + e.dst_type_id + ))); } let id = ts.nodes.len() as u32; let node = NodeRecord { @@ -3169,7 +3176,15 @@ impl DiskGraphEngine { } else { r }; - self.emit_edge(layer_idx, e.src_local_id, dst_local_id, e.rel_id, e.swap_edge_direction, state, output)?; + self.emit_edge( + layer_idx, + e.src_local_id, + dst_local_id, + e.rel_id, + e.swap_edge_direction, + state, + output, + )?; } let mut lo = 0; @@ -3253,7 +3268,15 @@ impl DiskGraphEngine { }; // Apply edge direction reversal if requested (PyG CSC convention). - self.emit_edge(layer_idx, src_local_id, dst_local_id, rel_id, swap_edge_direction, state, output) + self.emit_edge( + layer_idx, + src_local_id, + dst_local_id, + rel_id, + swap_edge_direction, + state, + output, + ) } } diff --git a/packages/diskgraph/diskgraph-engine/src/state.rs b/packages/diskgraph/diskgraph-engine/src/state.rs index 0ea0a261d..bd351be37 100644 --- a/packages/diskgraph/diskgraph-engine/src/state.rs +++ b/packages/diskgraph/diskgraph-engine/src/state.rs @@ -429,12 +429,16 @@ mod tests { let pool = StatePool::new(); { let mut state = pool.acquire(); - state.raw_edges = (0..4).map(|_| Vec::with_capacity(RESET_THRESHOLD * 2)).collect(); + state.raw_edges = (0..4) + .map(|_| Vec::with_capacity(RESET_THRESHOLD * 2)) + .collect(); let cs = &mut state.commit_scratch; cs.order.reserve(RESET_THRESHOLD * 4); cs.res.reserve(RESET_THRESHOLD * 4); cs.new_edges.reserve(RESET_THRESHOLD * 4); - state.seed_maps = (0..64).map(|_| GxHashMap::with_capacity(RESET_THRESHOLD / 16)).collect(); + state.seed_maps = (0..64) + .map(|_| GxHashMap::with_capacity(RESET_THRESHOLD / 16)) + .collect(); } let state = pool.acquire(); let edge_capacity: usize = state.raw_edges.iter().map(Vec::capacity).sum(); diff --git a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs index 969599e45..a31e09f0d 100644 --- a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs +++ b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs @@ -1448,7 +1448,11 @@ fn write_adjacency_shard( for &len in &lens { offsets.extend_from_slice(&(len as u16).to_le_bytes()); } - assert_eq!(lens.iter().sum::(), edges_len, "offsets must sum to edge file size"); + assert_eq!( + lens.iter().sum::(), + edges_len, + "offsets must sum to edge file size" + ); let mut fo = File::create(dir.join(OFFSETS_FILENAME)).unwrap(); fo.write_all(&offsets).unwrap(); @@ -2505,7 +2509,12 @@ fn test_finalize_preserves_anchor_major_seed_prefix() { ); } -type SampleSnapshot = (Vec<(usize, u16, u32, u32)>, Vec>, Vec>, Vec>); +type SampleSnapshot = ( + Vec<(usize, u16, u32, u32)>, + Vec>, + Vec>, + Vec>, +); fn random_adjacency(rng: &mut rand_chacha::ChaCha8Rng, n: u8) -> Vec> { use rand::Rng; @@ -2513,7 +2522,12 @@ fn random_adjacency(rng: &mut rand_chacha::ChaCha8Rng, n: u8) -> Vec = (0..rng.gen_range(0..30)) - .map(|_| (rng.gen_range(0..n), sk.encode_lower(rng.gen_range(0..1_000)))) + .map(|_| { + ( + rng.gen_range(0..n), + sk.encode_lower(rng.gen_range(0..1_000)), + ) + }) .collect(); e.sort_by_key(|&(_, k)| std::cmp::Reverse(k)); e @@ -2529,13 +2543,35 @@ fn per_seed_request( i: usize, ) -> BatchSamplingRequest { use rand::Rng; - let mode = if i % 2 == 0 { SamplingMode::TemporalMostRecent } else { SamplingMode::Uniform }; + let mode = if i % 2 == 0 { + SamplingMode::TemporalMostRecent + } else { + SamplingMode::Uniform + }; let mut req = make_sampling_request_with_mode(vec![], vec![6, 4, 3, 5], mode); for layer in &mut req.plan.layers { layer.default_fanout = layer.fanouts.drain().next().map(|(_, f)| f); } - req.seed_nodes = types.iter().map(|t| (t.to_string(), (0..seeds).map(|_| rng.gen_range(0..n as u64)).collect())).collect(); - req.seed_timestamps = Some(types.iter().map(|t| (t.to_string(), (0..seeds).map(|_| rng.gen_range(0..1_000)).collect())).collect()); + req.seed_nodes = types + .iter() + .map(|t| { + ( + t.to_string(), + (0..seeds).map(|_| rng.gen_range(0..n as u64)).collect(), + ) + }) + .collect(); + req.seed_timestamps = Some( + types + .iter() + .map(|t| { + ( + t.to_string(), + (0..seeds).map(|_| rng.gen_range(0..1_000)).collect(), + ) + }) + .collect(), + ); req.seed_timestamp_unit = diskgraph_core::time::TimeUnit::Seconds; req.flags.swap_edge_direction = i % 3 == 0; req @@ -2543,22 +2579,40 @@ fn per_seed_request( /// Samples every request on both recording paths and asserts identical /// edges, node records and stats. -fn assert_per_seed_matches_legacy(sampler: &mut DiskGraphEngine, requests: &[BatchSamplingRequest]) -> Vec { +fn assert_per_seed_matches_legacy( + sampler: &mut DiskGraphEngine, + requests: &[BatchSamplingRequest], +) -> Vec { let run = |sampler: &DiskGraphEngine| -> Vec { let pool = crate::state::StatePool::new(); requests .iter() .map(|req| { let mut state = pool.acquire(); - let mut output = MockOutput { edges: vec![], feature_batches: vec![], schema: None, fused_batches: vec![] }; + let mut output = MockOutput { + edges: vec![], + feature_batches: vec![], + schema: None, + fused_batches: vec![], + }; sampler.sample(req, &mut state, &mut output).unwrap(); let mut types: Vec<_> = state.type_states.iter().collect(); types.sort_by_key(|(t, _)| **t); let nodes = types .iter() - .map(|(_, ts)| ts.nodes.iter().map(|r| (r.global_id, r.seed_index, r.local_id, r.seed_type_id)).collect()) + .map(|(_, ts)| { + ts.nodes + .iter() + .map(|r| (r.global_id, r.seed_index, r.local_id, r.seed_type_id)) + .collect() + }) .collect(); - (output.edges, nodes, state.stats.nodes_per_type.clone(), state.stats.edges_per_type.clone()) + ( + output.edges, + nodes, + state.stats.nodes_per_type.clone(), + state.stats.edges_per_type.clone(), + ) }) .collect() }; @@ -2581,8 +2635,24 @@ fn per_seed_recording_matches_legacy_recording() { let (reader, _) = write_adjacency_shard(dir.path(), &adjacency, 16, &node_data); let reader = Arc::new(reader); let mut sampler = DiskGraphEngine::new( - [("User".to_string(), vec![ShardHandle { reader, min_id: 0, max_id: n as u64 - 1 }])].into(), - [(0, ("Connects".to_string(), "User".to_string(), "User".to_string()))].into(), + [( + "User".to_string(), + vec![ShardHandle { + reader, + min_id: 0, + max_id: n as u64 - 1, + }], + )] + .into(), + [( + 0, + ( + "Connects".to_string(), + "User".to_string(), + "User".to_string(), + ), + )] + .into(), [(0usize, SortKeyDecoder::new(&temporal_sortkey_meta()))].into(), [("User".to_string(), RowPackDecoder::new(ctx).unwrap())].into(), None, @@ -2591,10 +2661,13 @@ fn per_seed_recording_matches_legacy_recording() { Some(4), ); - let requests: Vec = - (0..6).map(|i| per_seed_request(&mut rng, &["User"], n, [1, 40, 700][i % 3], i)).collect(); + let requests: Vec = (0..6) + .map(|i| per_seed_request(&mut rng, &["User"], n, [1, 40, 700][i % 3], i)) + .collect(); let per_seed = assert_per_seed_matches_legacy(&mut sampler, &requests); - let (edges, nodes) = per_seed.iter().fold((0, 0), |(e, v), s| (e + s.0.len(), v + s.1[0].len())); + let (edges, nodes) = per_seed + .iter() + .fold((0, 0), |(e, v), s| (e + s.0.len(), v + s.1[0].len())); assert!(edges > nodes, "the graph must revisit nodes within a seed"); } @@ -2607,12 +2680,23 @@ fn per_seed_recording_matches_legacy_recording_across_seed_types() { let dirs = [TempDir::new().unwrap(), TempDir::new().unwrap()]; let readers: Vec> = dirs .iter() - .map(|dir| Arc::new(write_adjacency_shard(dir.path(), &random_adjacency(&mut rng, n), 16, &node_data).0)) + .map(|dir| { + Arc::new( + write_adjacency_shard(dir.path(), &random_adjacency(&mut rng, n), 16, &node_data).0, + ) + }) + .collect(); + let requests: Vec = (0..6) + .map(|i| per_seed_request(&mut rng, &["A", "B"], n, [3, 40, 700][i % 3], i)) .collect(); - let requests: Vec = - (0..6).map(|i| per_seed_request(&mut rng, &["A", "B"], n, [3, 40, 700][i % 3], i)).collect(); for threads in [1, 4] { - let shard = |t: usize| vec![ShardHandle { reader: readers[t].clone(), min_id: 0, max_id: n as u64 - 1 }]; + let shard = |t: usize| { + vec![ShardHandle { + reader: readers[t].clone(), + min_id: 0, + max_id: n as u64 - 1, + }] + }; let mut sampler = DiskGraphEngine::new( [("A".to_string(), shard(0)), ("B".to_string(), shard(1))].into(), [ @@ -2620,7 +2704,11 @@ fn per_seed_recording_matches_legacy_recording_across_seed_types() { (1, ("BToA".to_string(), "B".to_string(), "A".to_string())), ] .into(), - [(0usize, SortKeyDecoder::new(&temporal_sortkey_meta())), (1usize, SortKeyDecoder::new(&temporal_sortkey_meta()))].into(), + [ + (0usize, SortKeyDecoder::new(&temporal_sortkey_meta())), + (1usize, SortKeyDecoder::new(&temporal_sortkey_meta())), + ] + .into(), [ ("A".to_string(), RowPackDecoder::new(ctx.clone()).unwrap()), ("B".to_string(), RowPackDecoder::new(ctx.clone()).unwrap()), @@ -2632,7 +2720,9 @@ fn per_seed_recording_matches_legacy_recording_across_seed_types() { Some(threads), ); let per_seed = assert_per_seed_matches_legacy(&mut sampler, &requests); - assert!(per_seed.iter().all(|s| s.1.len() == 2 && s.0.iter().any(|e| e.1 == 1))); + assert!(per_seed + .iter() + .all(|s| s.1.len() == 2 && s.0.iter().any(|e| e.1 == 1))); } } @@ -2673,7 +2763,12 @@ fn repeated_runs_give_identical_output() { let runs: Vec<_> = (0..16) .map(|_| { let mut state = crate::state::StatePool::new().acquire(); - let mut output = MockOutput { edges: vec![], feature_batches: vec![], schema: None, fused_batches: vec![] }; + let mut output = MockOutput { + edges: vec![], + feature_batches: vec![], + schema: None, + fused_batches: vec![], + }; sampler.sample(&req, &mut state, &mut output).unwrap(); let rhs: Vec<_> = state.type_states[&1].nodes.iter().map(|n| (n.global_id, n.seed_type_id)).collect(); (output.edges, rhs, output.fused_batches) From 0d8c83252854239648f397fb1c3959a129d6fb0b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 22:05:38 +0000 Subject: [PATCH 9/9] diskgraph: use is_multiple_of in the per-seed test request builder Clippy's manual_is_multiple_of fires once the request index is a usize. Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-engine/src/tests_unit.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs index a31e09f0d..9f416b487 100644 --- a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs +++ b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs @@ -2543,7 +2543,7 @@ fn per_seed_request( i: usize, ) -> BatchSamplingRequest { use rand::Rng; - let mode = if i % 2 == 0 { + let mode = if i.is_multiple_of(2) { SamplingMode::TemporalMostRecent } else { SamplingMode::Uniform @@ -2573,7 +2573,7 @@ fn per_seed_request( .collect(), ); req.seed_timestamp_unit = diskgraph_core::time::TimeUnit::Seconds; - req.flags.swap_edge_direction = i % 3 == 0; + req.flags.swap_edge_direction = i.is_multiple_of(3); req }