use super::core::EdgeStore;
use crate::edge::csr_trait::MutableCsrTrait;
use crate::edge::{CsrShardSet, EdgeStrategy, Nbr, RecordForm};
use linkrs_core::types::{Timestamp, VertexId};
use linkrs_core::StorageResult;
use std::collections::HashMap;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RemapStats {
pub row_misses: usize,
pub neighbor_misses: usize,
pub entries: usize,
}
fn remapped_row_counted(id: u32, mapping: Option<&HashMap<u32, u32>>, misses: &mut usize) -> u32 {
match mapping {
Some(m) => match m.get(&id).copied() {
Some(new_id) => new_id,
None => {
*misses += 1;
id
}
},
None => id,
}
}
fn remap_endpoint_key_counted(
key: VertexId,
mapping: Option<&HashMap<u32, u32>>,
misses: &mut usize,
) -> VertexId {
let Some((endpoint, rank)) = EdgeStore::try_decode_edge_endpoint(key) else {
*misses += 1;
return key;
};
match endpoint.as_internal_u32() {
Some(id) => match mapping {
Some(m) => match m.get(&id).copied() {
Some(new_id) => EdgeStore::edge_endpoint_key(new_id, rank),
None => {
*misses += 1;
key
}
},
None => key,
},
None => {
*misses += 1;
key
}
}
}
#[derive(Clone, Copy)]
struct RebuildLayout {
strategy: EdgeStrategy,
group_bits: u32,
overflow_chunk_edges: usize,
form: RecordForm,
}
fn remap_direction(
old: &CsrShardSet,
row_mapping: Option<&HashMap<u32, u32>>,
neighbor_mapping: Option<&HashMap<u32, u32>>,
layout: RebuildLayout,
stats: &mut RemapStats,
) -> StorageResult<CsrShardSet> {
let RebuildLayout {
strategy,
group_bits,
overflow_chunk_edges,
form,
} = layout;
let mut rebuilt = CsrShardSet::new(strategy, group_bits, overflow_chunk_edges, form)?;
if strategy == EdgeStrategy::None {
return Ok(rebuilt);
}
for gid in old.existing_group_ids() {
let entries: Vec<(u32, Nbr)> = old
.group_variant(gid)
.map(|variant| {
variant
.iter_all()
.map(|(src, nbr)| {
let base = crate::edge::node_group::group_base(gid, old.group_bits());
let global = src.as_internal_u32().unwrap_or(0).saturating_add(base);
(global, nbr)
})
.collect()
})
.unwrap_or_default();
for (src, nbr) in entries {
let mut row_miss = 0usize;
let mut nbr_miss = 0usize;
let new_src = remapped_row_counted(src, row_mapping, &mut row_miss);
let new_neighbor =
remap_endpoint_key_counted(nbr.to_vertex_id(), neighbor_mapping, &mut nbr_miss);
let new_nbr =
match new_neighbor
.try_decode_edge_endpoint()
.and_then(|(ep_vid, ep_rank)| {
ep_vid.as_internal_u32().map(|endpoint| (endpoint, ep_rank))
}) {
Some((endpoint, rank)) => Nbr {
endpoint,
rank,
..nbr
},
None => {
nbr_miss += 1;
nbr
}
};
stats.row_misses += row_miss;
stats.neighbor_misses += nbr_miss;
stats.entries += 1;
let inline_value = if form == RecordForm::Bundled {
let base = crate::edge::node_group::group_base(gid, old.group_bits());
let old_local = src.saturating_sub(base);
old.group_variant(gid)
.and_then(|variant| variant.bundled_value_by_edge_id(old_local, nbr.edge_id))
.and_then(|(raw, valid)| valid.then_some(raw))
} else {
None
};
if form == RecordForm::Bundled {
rebuilt.bundled_insert_with_value(
new_src,
new_nbr.to_vertex_id(),
new_nbr.edge_id,
0,
inline_value,
)?;
} else {
rebuilt.insert_edge(new_src, new_nbr.to_vertex_id(), new_nbr.edge_id, 0)?;
}
if new_nbr.delete_ts != Timestamp::MAX {
let _ = rebuilt.delete_edge(new_src, new_nbr.edge_id, new_nbr.delete_ts);
rebuilt.mark_reclaim_hint_for(new_src);
}
}
}
rebuilt.mark_all_dirty();
rebuilt.clear_all_append_logs();
rebuilt.truncate_trailing_empty_groups();
Ok(rebuilt)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ReshardStats {
pub old_bits: u32,
pub new_bits: u32,
pub out_groups: usize,
pub in_groups: usize,
pub edges: u64,
}
impl EdgeStore {
pub fn reshard(&mut self, new_group_bits: u32) -> StorageResult<ReshardStats> {
crate::edge::node_group::validate_group_bits(new_group_bits)?;
let old_bits = self.config.node_group_bits;
if new_group_bits == old_bits {
return Ok(ReshardStats {
old_bits,
new_bits: new_group_bits,
out_groups: self.out_csr.group_count(),
in_groups: self.in_csr.group_count(),
edges: self.edge_count(),
});
}
let mut stats = RemapStats::default();
self.out_csr = remap_direction(
&self.out_csr,
None,
None,
RebuildLayout {
strategy: self.schema.oe_strategy,
group_bits: new_group_bits,
overflow_chunk_edges: self.config.overflow_chunk_edges,
form: self.out_csr.record_form(),
},
&mut stats,
)?;
self.in_csr = remap_direction(
&self.in_csr,
None,
None,
RebuildLayout {
strategy: self.schema.ie_strategy,
group_bits: new_group_bits,
overflow_chunk_edges: self.config.overflow_chunk_edges,
form: self.in_csr.record_form(),
},
&mut stats,
)?;
self.config.node_group_bits = new_group_bits;
let owner_stats = self.rebuild_owner_map_with_stats();
if let Some(stats) = &self.stats_manager {
stats.add_value_with_amount(
linkrs_metrics::MetricType::EdgeRelocatedOrphans,
owner_stats.relocated_orphans as u64,
);
}
let drift = self.audit_copy_drift();
if !drift.is_empty() {
return Err(linkrs_core::StorageError::db_error(format!(
"reshard drifted authority and projection: {}",
drift.join("; ")
)));
}
self.migration_pending_checkpoint = true;
log::debug!(
"EdgeTable[{}] resharded width {} -> {}; out_groups={}, in_groups={}, edges={}",
self.label,
old_bits,
new_group_bits,
self.out_csr.group_count(),
self.in_csr.group_count(),
self.edge_count(),
);
Ok(ReshardStats {
old_bits,
new_bits: new_group_bits,
out_groups: self.out_csr.group_count(),
in_groups: self.in_csr.group_count(),
edges: self.edge_count(),
})
}
pub fn remap_vertex_ids(
&mut self,
src_mapping: Option<&HashMap<u32, u32>>,
dst_mapping: Option<&HashMap<u32, u32>>,
) -> StorageResult<()> {
self.remap_vertex_ids_with_stats(src_mapping, dst_mapping)
.map(|_| ())
}
pub fn remap_vertex_ids_with_stats(
&mut self,
src_mapping: Option<&HashMap<u32, u32>>,
dst_mapping: Option<&HashMap<u32, u32>>,
) -> StorageResult<RemapStats> {
if src_mapping.is_none() && dst_mapping.is_none() {
return Ok(RemapStats::default());
}
let src_empty = src_mapping.is_none_or(|m| m.is_empty());
let dst_empty = dst_mapping.is_none_or(|m| m.is_empty());
if src_empty && dst_empty {
return Ok(RemapStats::default());
}
let mut stats = RemapStats::default();
self.out_csr = remap_direction(
&self.out_csr,
src_mapping,
dst_mapping,
RebuildLayout {
strategy: self.schema.oe_strategy,
group_bits: self.config.node_group_bits,
overflow_chunk_edges: self.config.overflow_chunk_edges,
form: self.out_csr.record_form(),
},
&mut stats,
)?;
self.in_csr = remap_direction(
&self.in_csr,
dst_mapping,
src_mapping,
RebuildLayout {
strategy: self.schema.ie_strategy,
group_bits: self.config.node_group_bits,
overflow_chunk_edges: self.config.overflow_chunk_edges,
form: self.in_csr.record_form(),
},
&mut stats,
)?;
if self.property_index.is_some() {
let pool_capacity = self
.property_index
.as_ref()
.map(|idx| idx.pool_capacity())
.unwrap_or(1024);
self.build_property_index(pool_capacity)?;
}
let owner_stats = self.rebuild_owner_map_with_stats();
if let Some(metrics) = &self.stats_manager {
metrics.add_value_with_amount(
linkrs_metrics::MetricType::EdgeRelocatedOrphans,
owner_stats.relocated_orphans as u64,
);
}
let drift = self.audit_copy_drift();
if !drift.is_empty() {
return Err(linkrs_core::StorageError::db_error(format!(
"remap drifted authority and projection: {}",
drift.join("; ")
)));
}
let (_, _, live_orphans) = self.copy_audit();
if let Some(stats) = &self.stats_manager {
stats.add_value_with_amount(
linkrs_metrics::MetricType::EdgeLiveAuthorityOrphans,
live_orphans as u64,
);
}
log::debug!(
"EdgeTable[{}] remapped vertex IDs (src_mapping={}, dst_mapping={}); out_groups={}, in_groups={}, row_misses={}, neighbor_misses={}, entries={}",
self.label,
src_mapping.map(|m| m.len()).unwrap_or(0),
dst_mapping.map(|m| m.len()).unwrap_or(0),
self.out_csr.group_count(),
self.in_csr.group_count(),
stats.row_misses,
stats.neighbor_misses,
stats.entries,
);
Ok(stats)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::edge::edge_table::config::EdgeTableConfig;
use crate::edge::{EdgeSchema, EdgeStrategy, RecordForm};
use crate::types::StoragePropertyDef;
use linkrs_core::Value;
fn make_table() -> EdgeStore {
let schema = EdgeSchema {
label_id: 0,
label_name: "knows".into(),
src_label: 0,
dst_label: 0,
properties: vec![StoragePropertyDef {
name: "weight".into(),
data_type: linkrs_core::types::DataType::Double,
nullable: false,
default_value: Some(Value::Double(0.0)),
}],
oe_strategy: EdgeStrategy::Multiple,
ie_strategy: EdgeStrategy::Multiple,
schema_version: 1,
record_form: RecordForm::default(),
};
EdgeStore::with_config(schema, EdgeTableConfig::default()).unwrap()
}
fn mapping_from_removals(live: &[u32]) -> HashMap<u32, u32> {
let mut mapping = HashMap::new();
for (new_id, &old_id) in live.iter().enumerate() {
if old_id != new_id as u32 {
mapping.insert(old_id, new_id as u32);
}
}
mapping
}
fn make_single_table() -> EdgeStore {
let schema = EdgeSchema {
label_id: 0,
label_name: "single".into(),
src_label: 0,
dst_label: 0,
properties: vec![],
oe_strategy: EdgeStrategy::Single,
ie_strategy: EdgeStrategy::Single,
schema_version: 1,
record_form: RecordForm::default(),
};
EdgeStore::with_config(schema, EdgeTableConfig::default()).unwrap()
}
#[test]
fn test_remap_rows_and_neighbors() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
table.insert_edge(1, 3, 1, &[], 100).unwrap();
table.insert_edge(3, 5, 2, &[], 100).unwrap();
table.insert_edge(5, 0, 3, &[], 100).unwrap();
let mapping = mapping_from_removals(&[0, 1, 3, 5]);
table
.remap_vertex_ids(Some(&mapping), Some(&mapping))
.expect("remap should succeed");
assert_eq!(table.out_edges(0, 200).len(), 1);
assert_eq!(table.out_edges(1, 200).len(), 1);
assert_eq!(table.out_edges(2, 200).len(), 1);
assert_eq!(table.out_edges(3, 200).len(), 1);
assert!(table.get_edge(0, 1, 0, 200).is_some());
assert!(table.get_edge(1, 2, 1, 200).is_some());
assert!(table.get_edge(2, 3, 2, 200).is_some());
assert!(table.get_edge(3, 0, 3, 200).is_some());
assert_eq!(table.out_csr.group_count(), 1);
assert_eq!(table.in_csr.group_count(), 1);
assert!(table.get_edge(4, 0, 0, 200).is_none());
assert_eq!(table.out_edges(4, 200).len(), 0);
}
#[test]
fn test_remap_preserves_tombstoned_entries() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
table.insert_edge(3, 4, 1, &[], 100).unwrap();
table.delete_edge(0, 1, 0, 150).unwrap();
let mapping = mapping_from_removals(&[0, 1, 3, 4]);
table
.remap_vertex_ids(Some(&mapping), Some(&mapping))
.unwrap();
let before = table.out_edges(0, 149);
assert_eq!(before.len(), 1);
assert_eq!(before[0].dst_vid.as_int64(), Some(1));
assert_eq!(table.out_edges(0, 200).len(), 0);
assert!(table.get_edge(2, 3, 1, 200).is_some());
assert_eq!(table.out_csr.group_count(), 1);
}
#[test]
fn test_remap_dst_only_affects_neighbors_out() {
let mut table = make_table();
table.insert_edge(1, 4, 0, &[], 100).unwrap();
let mapping = HashMap::from([(4u32, 0u32)]);
table.remap_vertex_ids(None, Some(&mapping)).unwrap();
assert!(table.get_edge(1, 0, 0, 200).is_some());
assert!(table.get_edge(1, 4, 0, 200).is_none());
assert_eq!(table.out_csr.group_count(), 1);
}
#[test]
fn test_remap_src_only_affects_rows() {
let mut table = make_table();
table.insert_edge(4, 1, 0, &[], 100).unwrap();
let mapping = HashMap::from([(4u32, 0u32)]);
table.remap_vertex_ids(Some(&mapping), None).unwrap();
assert!(table.get_edge(0, 1, 0, 200).is_some());
assert!(table.get_edge(4, 1, 0, 200).is_none());
assert_eq!(table.out_csr.group_count(), 1);
assert_eq!(table.in_csr.group_count(), 1);
}
#[test]
fn test_remap_empty_mapping_is_noop() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
let groups_before = table.out_csr.group_count();
table.remap_vertex_ids(Some(&HashMap::new()), None).unwrap();
assert_eq!(table.out_csr.group_count(), groups_before);
assert!(table.get_edge(0, 1, 0, 200).is_some());
}
#[test]
fn test_remap_moves_edges_across_groups() {
let mut table = make_table();
table.insert_edge(5000, 6000, 0, &[], 100).unwrap();
assert_eq!(table.out_csr.group_count(), 2);
let src_mapping = HashMap::from([(5000u32, 1u32)]);
let dst_mapping = HashMap::from([(6000u32, 2u32)]);
table
.remap_vertex_ids(Some(&src_mapping), Some(&dst_mapping))
.unwrap();
assert!(table.get_edge(1, 2, 0, 200).is_some());
assert!(table.get_edge(5000, 6000, 0, 200).is_none());
assert_eq!(table.out_csr.group_count(), 1);
}
#[test]
fn test_remap_flush_load_roundtrip() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
table.insert_edge(5000, 6000, 0, &[], 100).unwrap();
let dir = tempfile::tempdir().expect("temporary edge table directory");
table
.flush(
dir.path(),
crate::compression::CompressionType::Zstd { level: 3 },
)
.expect("baseline flush should succeed");
let src_mapping = HashMap::from([(5000u32, 1u32)]);
let dst_mapping = HashMap::from([(6000u32, 2u32)]);
table
.remap_vertex_ids(Some(&src_mapping), Some(&dst_mapping))
.unwrap();
assert!(!table.out_csr.dirty_group_ids().is_empty());
table
.flush(
dir.path(),
crate::compression::CompressionType::Zstd { level: 3 },
)
.expect("post-remap flush should succeed");
let mut loaded = make_table();
loaded.load(dir.path()).expect("load should succeed");
assert!(loaded.get_edge(0, 1, 0, 200).is_some());
assert!(loaded.get_edge(1, 2, 0, 200).is_some());
assert!(loaded.get_edge(5000, 6000, 0, 200).is_none());
assert_eq!(loaded.edge_count(), 2);
}
#[test]
fn test_remap_compaction_reports_zero_misses() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
table.insert_edge(1, 0, 0, &[], 100).unwrap();
let mapping = HashMap::from([(0u32, 0u32), (1u32, 1u32)]);
let stats = table
.remap_vertex_ids_with_stats(Some(&mapping), Some(&mapping))
.unwrap();
assert_eq!(stats.row_misses, 0);
assert_eq!(stats.neighbor_misses, 0);
assert!(table.get_edge(0, 1, 0, 200).is_some());
assert!(table.get_edge(1, 0, 0, 200).is_some());
assert_eq!(table.live_authority_orphans(), 0);
}
#[test]
fn test_remap_partial_mapping_reports_expected_misses() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
let mapping = HashMap::from([(99u32, 100u32)]);
let stats = table
.remap_vertex_ids_with_stats(Some(&mapping), Some(&mapping))
.unwrap();
assert!(stats.row_misses > 0);
assert!(stats.neighbor_misses > 0);
assert!(table.get_edge(0, 1, 0, 200).is_some());
}
#[test]
fn test_single_strategy_remap_into_fresh_shards() {
let mut table = make_single_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
let mapping = HashMap::from([(0u32, 0u32), (1u32, 1u32)]);
let stats = table
.remap_vertex_ids_with_stats(Some(&mapping), Some(&mapping))
.unwrap();
assert_eq!(stats.row_misses, 0);
assert_eq!(stats.neighbor_misses, 0);
assert!(table.get_edge(0, 1, 0, 200).is_some());
assert_eq!(table.live_authority_orphans(), 0);
}
#[test]
fn test_reshard_equal_width_is_noop_success_and_bad_width_rejected() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
let current = table.config.node_group_bits;
let stats = table.reshard(current).expect("equal width must succeed");
assert_eq!(stats.old_bits, current);
assert_eq!(stats.new_bits, current);
assert!(table.get_edge(0, 1, 0, 200).is_some());
assert!(table.audit_copy_drift().is_empty());
assert!(table.reshard(0).is_err());
assert!(table.reshard(21).is_err());
assert!(table.get_edge(0, 1, 0, 200).is_some());
}
#[test]
fn test_reshard_arms_mandatory_checkpoint_fence() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
let current = table.config.node_group_bits;
assert!(current > 1);
table.reshard(current - 1).expect("width change succeeds");
assert!(table.is_migration_checkpoint_required());
assert!(table.get_edge(0, 1, 0, 200).is_some());
assert!(table.insert_edge(2, 3, 0, &[], 200).is_err());
assert!(table.audit_copy_drift().is_empty());
let dir = tempfile::tempdir().expect("temporary edge table directory");
table
.flush(
dir.path(),
crate::compression::CompressionType::Zstd { level: 3 },
)
.expect("mandatory checkpoint");
assert!(!table.is_migration_checkpoint_required());
table
.insert_edge(2, 3, 0, &[], 200)
.expect("write after checkpoint");
assert!(table.get_edge(2, 3, 0, 300).is_some());
assert!(table.audit_copy_drift().is_empty());
}
#[test]
fn test_edge_id_counter_monotonic_across_reload() {
let mut table = make_table();
table.insert_edge(0, 1, 0, &[], 100).unwrap();
table.insert_edge(0, 2, 0, &[], 110).unwrap();
table.delete_edge(0, 2, 0, 150).unwrap();
let before = table.next_edge_id;
let dir = tempfile::tempdir().expect("temporary edge table directory");
table
.flush(
dir.path(),
crate::compression::CompressionType::Zstd { level: 3 },
)
.expect("flush should succeed");
let mut loaded = make_table();
loaded.load(dir.path()).expect("load should succeed");
assert!(loaded.next_edge_id >= before);
let mut out_seen = std::collections::HashSet::new();
for (_, nbr) in loaded.out_csr.iter_all() {
assert!(out_seen.insert(nbr.edge_id));
}
let mut in_seen = std::collections::HashSet::new();
for (_, nbr) in loaded.in_csr.iter_all() {
assert!(in_seen.insert(nbr.edge_id));
}
assert_eq!(out_seen, in_seen);
}
}