1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
use std::{
collections::HashMap,
path::{Path, PathBuf},
sync::Arc,
};
use arrow_array::{Int64Array, RecordBatch};
use arrow_schema::{DataType as ArrowDataType, Field, Schema};
use crate::{
error::Result,
io,
metadata::{EdgeInfo, PropertyGroup},
types::AdjListType,
writer::vertex_writer::{PropertyValue, build_record_batch_for_edges},
};
/// An edge with source/destination and optional properties.
#[derive(Debug, Clone)]
pub struct Edge {
pub src: i64,
pub dst: i64,
properties: HashMap<String, PropertyValue>,
}
impl Edge {
pub fn new(src: i64, dst: i64) -> Self {
Self {
src,
dst,
properties: HashMap::new(),
}
}
pub fn add_property(mut self, name: impl Into<String>, value: PropertyValue) -> Self {
self.properties.insert(name.into(), value);
self
}
pub fn get(&self, name: &str) -> Option<&PropertyValue> {
self.properties.get(name)
}
}
/// Builds and writes edge chunks (adj list + properties + offsets) to disk.
pub struct EdgesBuilder {
info: EdgeInfo,
base_dir: PathBuf,
adj_type: AdjListType,
/// Total number of vertices in the aligned dimension (src or dst).
vertex_count: u64,
edges: Vec<Edge>,
}
impl EdgesBuilder {
pub fn new(
info: EdgeInfo,
base_dir: impl AsRef<Path>,
adj_type: AdjListType,
vertex_count: u64,
) -> Self {
Self {
info,
base_dir: base_dir.as_ref().to_path_buf(),
adj_type,
vertex_count,
edges: Vec::new(),
}
}
pub fn add_edge(&mut self, edge: Edge) {
self.edges.push(edge);
}
/// Sort edges by the alignment key and flush as chunked files.
pub fn dump(&mut self) -> Result<()> {
// Chunk membership tests the alignment id as a `u64` (see `key` below),
// so sort by that same `u64` key. For the only ids that ever land in a
// chunk — non-negative vertex ids in `[0, vertex_count)` — this orders
// identically to sorting by the raw `i64`; out-of-range ids (negatives,
// which cast to huge `u64`) are dropped either way, so the written
// output is unchanged. Keeping the sort key equal to the membership key
// is what lets the single advancing cursor below stay correct.
let aligned_by_src = self.adj_type.aligned_by_src();
let key = |e: &Edge| -> u64 {
if aligned_by_src {
e.src as u64
} else {
e.dst as u64
}
};
self.edges.sort_unstable_by_key(key);
let vchunk_size = if aligned_by_src {
self.info.src_chunk_size
} else {
self.info.dst_chunk_size
};
let num_vertex_chunks = self.vertex_count.div_ceil(vchunk_size);
let file_type = self
.info
.get_adj_list_info(&self.adj_type)
.map(|a| a.file_type.clone())
.unwrap_or_default();
// The edges are now sorted ascending by the alignment key, and the
// vertex chunks below are visited with monotonically increasing
// `[v_start, v_end)` ranges. So a single cursor can walk the sorted Vec
// once, yielding the contiguous slice for each chunk in O(edges) total
// instead of re-scanning every edge per vertex chunk.
//
// The first chunk starts at key 0 and the cursor begins at the front.
// Edges whose key lands beyond the final chunk's `v_end`
// (>= vertex_count) are simply never reached and thus dropped, matching
// the original per-chunk filter.
let mut cursor = 0usize;
for vc in 0..num_vertex_chunks {
let v_start = vc * vchunk_size;
let v_end = ((vc + 1) * vchunk_size).min(self.vertex_count);
// Edges belonging to this vertex chunk: a contiguous run in the
// sorted Vec starting at `cursor`. Advance `cursor` to the first
// edge with `key >= v_end`.
debug_assert!(self.edges[cursor..].iter().all(|e| key(e) >= v_start));
let chunk_end = cursor + self.edges[cursor..].partition_point(|e| key(e) < v_end);
let chunk_edges: Vec<&Edge> = self.edges[cursor..chunk_end].iter().collect();
cursor = chunk_end;
// Write adjacency list chunks.
let edge_chunk_size = self.info.chunk_size as usize;
for (ec, edge_slice) in chunk_edges.chunks(edge_chunk_size).enumerate() {
let batch = build_adj_list_batch(edge_slice)?;
let path = self.base_dir.join(self.info.adj_list_chunk_path(
&self.adj_type,
vc,
ec as u64,
)?);
io::write_chunk(&path, &[batch], &file_type)?;
// Write property group chunks.
for group in self.info.property_groups.iter() {
let prop_batch = build_edge_prop_batch(edge_slice, group)?;
let prop_path = self.base_dir.join(self.info.property_chunk_path(
&self.adj_type,
group,
vc,
ec as u64,
));
io::write_chunk(&prop_path, &[prop_batch], &group.file_type)?;
}
}
// Write offset chunk for ordered adj lists.
if self.adj_type.is_ordered() {
let offset_batch =
build_offset_batch(&chunk_edges, v_start, v_end, vchunk_size, aligned_by_src)?;
let offset_path = self
.base_dir
.join(self.info.offset_chunk_path(&self.adj_type, vc)?);
io::write_chunk(&offset_path, &[offset_batch], &file_type)?;
}
}
self.edges.clear();
Ok(())
}
}
fn build_adj_list_batch(edges: &[&Edge]) -> Result<RecordBatch> {
let srcs: Vec<i64> = edges.iter().map(|e| e.src).collect();
let dsts: Vec<i64> = edges.iter().map(|e| e.dst).collect();
let schema = Arc::new(Schema::new(vec![
Field::new("src", ArrowDataType::Int64, false),
Field::new("dst", ArrowDataType::Int64, false),
]));
Ok(RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(srcs)),
Arc::new(Int64Array::from(dsts)),
],
)?)
}
fn build_offset_batch(
chunk_edges: &[&Edge],
v_start: u64,
v_end: u64,
_vchunk_size: u64,
aligned_by_src: bool,
) -> Result<RecordBatch> {
let count = (v_end - v_start) as usize + 1;
let mut offsets = vec![0i64; count];
// Histogram: count edges per vertex slot, then prefix-sum.
let mut hist = vec![0i64; (v_end - v_start) as usize];
for edge in chunk_edges {
// Key the histogram on the same vertex dimension the adj list is aligned
// by — `src` for *_by_source, `dst` for *_by_dest. Mirrors the `key`
// closure used in `dump()`; using `src` unconditionally would bucket
// ordered_by_dest offsets against the wrong dimension.
let key_id = if aligned_by_src { edge.src } else { edge.dst };
let slot = (key_id as u64).saturating_sub(v_start) as usize;
if slot < hist.len() {
hist[slot] += 1;
}
}
// Prefix sum -> offsets; first row is always 0.
let mut running = 0i64;
offsets[0] = 0;
for (i, count) in hist.iter().enumerate() {
running += count;
offsets[i + 1] = running;
}
let schema = Arc::new(Schema::new(vec![Field::new(
"offset",
ArrowDataType::Int64,
false,
)]));
Ok(RecordBatch::try_new(
schema,
vec![Arc::new(Int64Array::from(offsets))],
)?)
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::Array;
fn offsets_of(batch: &RecordBatch) -> Vec<i64> {
batch
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.to_vec()
}
#[test]
fn offset_batch_keys_on_dst_for_ordered_by_dest() {
// An OrderedByDest chunk is aligned/sorted by `dst`; the offset
// histogram must bucket on `dst`, not `src`.
//
// FAIL-ON-BUG: the old code keyed on `edge.src` unconditionally. Here
// every `src` (100) is outside [v_start, v_end) = [0, 3), so the buggy
// path drops all edges and yields all-zero offsets. PASS-ON-FIX: keying
// on `dst` produces the correct prefix sums.
let edges = vec![
Edge::new(100, 0),
Edge::new(100, 0),
Edge::new(100, 1),
Edge::new(100, 2),
];
let refs: Vec<&Edge> = edges.iter().collect();
// aligned_by_src = false → key on dst.
let batch = build_offset_batch(&refs, 0, 3, 4, false).unwrap();
assert_eq!(
offsets_of(&batch),
vec![0, 2, 3, 4],
"ordered_by_dest offsets must be the dst-keyed prefix sum",
);
// Sanity: keying on src (aligned_by_src = true) drops every edge here.
let buggy = build_offset_batch(&refs, 0, 3, 4, true).unwrap();
assert_eq!(offsets_of(&buggy), vec![0, 0, 0, 0]);
}
#[test]
fn offset_batch_keys_on_src_for_ordered_by_source() {
// The by-source path is unchanged: bucket on `src`.
let edges = vec![Edge::new(0, 9), Edge::new(1, 9), Edge::new(1, 9)];
let refs: Vec<&Edge> = edges.iter().collect();
let batch = build_offset_batch(&refs, 0, 2, 4, true).unwrap();
assert_eq!(offsets_of(&batch), vec![0, 1, 3]);
}
}
fn build_edge_prop_batch(edges: &[&Edge], group: &PropertyGroup) -> Result<RecordBatch> {
// Build pseudo-Vertex objects from edge properties for each property in the group.
let pseudo_vertices: Vec<crate::writer::vertex_writer::Vertex> = edges
.iter()
.map(|e| {
let mut v = crate::writer::vertex_writer::Vertex::new();
for prop in &group.properties {
if let Some(val) = e.get(&prop.name) {
v = v.add_property(prop.name.clone(), val.clone());
}
}
v
})
.collect();
build_record_batch_for_edges(&pseudo_vertices, group)
}