diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 6d91107fa..f537e6f91 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,24 @@ 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; + } + } + state.seed_slots = slots as usize; + if state.seed_maps.len() < state.seed_slots { + state + .seed_maps + .resize_with(state.seed_slots, 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 +2161,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 +2907,18 @@ 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() + .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 +2927,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 +3000,15 @@ 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; + state.raw_edges.extend(spare); + 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 +3024,8 @@ impl DiskGraphEngine { )?; } } + state.raw_edges = raw_edges_results; + state.raw_edges.extend(spare); } Ok(()) @@ -3054,6 +3095,134 @@ 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_slots; + + 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(); + 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, + 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 +3268,15 @@ 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..bd351be37 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) @@ -250,8 +251,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_slots: usize, + 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_MAPS_KEEP: usize = 1 << 10; + const RESET_THRESHOLD: usize = 1 << 18; // ~262k items impl Default for SamplingState { @@ -269,6 +309,11 @@ impl SamplingState { reservoir_scratch: Vec::with_capacity(128), 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(), } } @@ -281,6 +326,28 @@ impl SamplingState { self.key_scratch.clear(); 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); + } + 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); + let mut budget = RESET_THRESHOLD; + for m in &mut self.seed_maps { + m.clear(); + if m.capacity() > budget { + m.shrink_to(0); + } + budget -= m.capacity(); + } // Scratch buffer is inherently reusable } /// Access or create the state for a specific node type. @@ -357,6 +424,34 @@ 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); + 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(); + 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); + let map_capacity: usize = state.seed_maps.iter().map(GxHashMap::capacity).sum(); + assert!(map_capacity <= RESET_THRESHOLD); + assert!(map_capacity > 0); + } + #[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..9f416b487 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,46 @@ 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 +2509,223 @@ fn test_finalize_preserves_anchor_major_seed_prefix() { ); } +type SampleSnapshot = ( + Vec<(usize, u16, u32, u32)>, + Vec>, + Vec>, + Vec>, +); + +fn random_adjacency(rng: &mut rand_chacha::ChaCha8Rng, n: u8) -> Vec> { + use rand::Rng; + let sk = SortKeyEncoder::new(&temporal_sortkey_meta()); + (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() +} + +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.is_multiple_of(2) { + 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.is_multiple_of(3); + 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 + .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() + }; + sampler.legacy_recording = false; + let per_seed = run(sampler); + sampler.legacy_recording = true; + 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()); @@ -2534,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) 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"); + } +}