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
//! Bulk edge insert partitioned by group, with reservation pre-sizing.
use super::super::super::{
csr_shared::decode_endpoint_pair, CsrVariant, EdgePut, MutableCsrTrait, Nbr, RowEdgeBatch,
NO_EDGES_STORED_MSG,
};
use super::super::{group_id_for, local_vid};
use super::CsrShardSet;
use linkrs_core::types::{EdgeId, EdgeStrategy, Timestamp, VertexId};
use linkrs_core::{StorageError, StorageResult};
use std::collections::BTreeMap;
impl CsrShardSet {
/// Bulk insert pre-grouped edges with one reservation per touched row.
///
/// Input is `(src, dst, edge_id)` triples at global addresses sharing one
/// `ts`. The shared grouping and write path lives in
/// [`Self::batch_put_edges_with_ts`].
pub fn batch_put_edges(
&mut self,
edges: &[(u32, VertexId, EdgeId)],
ts: Timestamp,
check_duplicates: bool,
) -> StorageResult<usize> {
let quads: Vec<(u32, VertexId, EdgeId, Timestamp)> = edges
.iter()
.map(|&(src, dst, edge_id)| (src, dst, edge_id, ts))
.collect();
self.batch_put_edges_with_ts(&quads, check_duplicates)
}
/// Bulk insert with per-edge timestamps, partitioned by group.
///
/// Direct-write counterpart of `batch_put_edges` for initial loads with
/// heterogeneous timestamps: input is `(src, dst, edge_id, ts)` quads at
/// global addresses, rows are grouped by shard, each `Multiple` group is
/// written through its bulk path (single reservation, single live-set
/// rebuild per row) preserving each edge timestamp, and other variants
/// fall back to per-edge inserts. Dirt and append-log entries are
/// recorded per inserted edge exactly like the single-edge path.
/// Duplicate keys are rejected before any write when `check_duplicates`
/// is set; bulk-import callers prevalidate once and pass false.
pub fn batch_put_edges_with_ts(
&mut self,
edges: &[(u32, VertexId, EdgeId, Timestamp)],
check_duplicates: bool,
) -> StorageResult<usize> {
if self.strategy == EdgeStrategy::None {
return Err(StorageError::invalid_operation(
NO_EDGES_STORED_MSG.to_string(),
));
}
let mut by_group: BTreeMap<usize, Vec<(u32, VertexId, EdgeId, Timestamp)>> =
BTreeMap::new();
for (src, dst, edge_id, ts) in edges {
let gid = group_id_for(*src, self.group_bits);
by_group
.entry(gid)
.or_default()
.push((*src, *dst, *edge_id, *ts));
}
let mut inserted = 0usize;
for (gid, group_edges) in by_group {
let group_bits = self.group_bits;
self.ensure_group_id(gid)?;
let is_multiple = matches!(
self.shards.get(&gid).map(|shard| &shard.variant),
Some(CsrVariant::Multiple(_))
);
if is_multiple {
let batch: Vec<RowEdgeBatch> = {
let mut rows: BTreeMap<u32, Vec<EdgePut>> = BTreeMap::new();
for (src, dst, edge_id, ts) in &group_edges {
let local = local_vid(*src, group_bits);
let (endpoint, rank) = decode_endpoint_pair(*dst).ok_or_else(|| {
StorageError::invalid_input(format!(
"Malformed edge endpoint key: {}",
dst
))
})?;
rows.entry(local)
.or_default()
.push((endpoint, rank, *edge_id, *ts));
}
rows.into_iter().collect()
};
{
let shard = self.shards.get_mut(&gid).ok_or_else(|| {
StorageError::invalid_operation(format!("missing group {} on insert", gid))
})?;
let CsrVariant::Multiple(csr) = &mut shard.variant else {
return Err(StorageError::invalid_operation(format!(
"missing group {} on insert",
gid
)));
};
let counts: Vec<(u32, usize)> = batch
.iter()
.map(|(local, puts)| (*local, puts.len()))
.collect();
csr.reserve_for_batch(&counts);
inserted += csr.batch_put_edges(&batch, check_duplicates)?;
}
for (src, dst, edge_id, ts) in &group_edges {
let local = local_vid(*src, group_bits);
let (endpoint, rank) = decode_endpoint_pair(*dst).ok_or_else(|| {
StorageError::invalid_input(format!("Malformed edge endpoint key: {}", dst))
})?;
let nbr = Nbr::with_create_ts(endpoint, rank, *edge_id, *ts);
self.mark_region_insert(gid, local);
self.record_append_insert(gid, local, nbr);
}
} else {
for (src, dst, edge_id, ts) in group_edges {
let local = local_vid(src, group_bits);
{
let shard = self.shards.get_mut(&gid).ok_or_else(|| {
StorageError::invalid_operation(format!(
"missing group {} on insert",
gid
))
})?;
shard.variant.insert_edge(local, dst, edge_id, ts)?;
}
let (endpoint, rank) = decode_endpoint_pair(dst).ok_or_else(|| {
StorageError::invalid_input(format!("Malformed edge endpoint key: {}", dst))
})?;
let nbr = Nbr::with_create_ts(endpoint, rank, edge_id, ts);
self.mark_region_insert(gid, local);
self.record_append_insert(gid, local, nbr);
inserted += 1;
}
}
}
Ok(inserted)
}
/// Pre-size touched rows once for an incoming batch.
///
/// Groups missing rows are materialized first; each `Multiple` group
/// sizes its rows at the packed density target so the following
/// per-edge inserts land in reserved gaps. Other strategies skip
/// reservation. Only sizing happens here; inserts still flow through
/// the regular path so dirt and append logs stay exact.
pub fn reserve_for_batch(&mut self, counts: &[(u32, usize)]) -> StorageResult<()> {
if self.strategy == EdgeStrategy::None {
return Ok(());
}
let mut by_group: BTreeMap<usize, Vec<(u32, usize)>> = BTreeMap::new();
for (vid, incoming) in counts {
let gid = group_id_for(*vid, self.group_bits);
by_group
.entry(gid)
.or_default()
.push((local_vid(*vid, self.group_bits), *incoming));
}
for (gid, group_counts) in by_group {
self.ensure_group_id(gid)?;
let Some(shard) = self.shards.get_mut(&gid) else {
continue;
};
if let CsrVariant::Multiple(csr) = &mut shard.variant {
csr.reserve_for_batch(&group_counts);
}
}
Ok(())
}
}