1use crate::storage::delta::{ENTRY_SIZE_ESTIMATE, L1Entry, Op};
5use crate::storage::manager::StorageManager;
6use anyhow::{Result, anyhow};
7use arrow_array::Array;
8use arrow_array::builder::{ListBuilder, UInt64Builder};
9use arrow_array::{ListArray, RecordBatch, UInt64Array};
10use metrics;
11use std::collections::{HashMap, HashSet};
12use std::sync::Arc;
13use tracing::{error, info, instrument};
14use uni_common::core::id::{Eid, Vid};
15use uni_common::core::schema::DataType;
16use uni_common::{Properties, Value};
17use uni_crdt::Crdt;
18
19pub struct Compactor {
20 storage: Arc<StorageManager>,
21}
22
23impl Compactor {
24 pub fn new(storage: Arc<StorageManager>) -> Self {
25 Self { storage }
26 }
27
28 #[instrument(skip(self), level = "info")]
29 pub async fn compact_all(&self) -> Result<Vec<CompactionInfo>> {
30 let start = std::time::Instant::now();
31 let schema = self.storage.schema_manager().schema();
32 let mut compaction_results = Vec::new();
33
34 for label in schema.labels.keys() {
36 info!("Compacting vertices for label {}", label);
37 if let Err(e) = self.compact_vertices(label).await {
38 error!("Failed to compact vertices for {}: {}", label, e);
39 }
40 }
41
42 for (edge_type, meta) in &schema.edge_types {
44 for label in &meta.src_labels {
46 info!("Compacting adjacency {} -> {} (fwd)", label, edge_type);
47 match self.compact_adjacency(edge_type, label, "fwd").await {
48 Ok(info) => compaction_results.push(info),
49 Err(e) => {
50 error!(
51 "Failed to compact adjacency {} -> {}: {}",
52 label, edge_type, e
53 );
54 }
55 }
56 }
57
58 for label in &meta.dst_labels {
60 info!("Compacting adjacency {} <- {} (bwd)", label, edge_type);
61 match self.compact_adjacency(edge_type, label, "bwd").await {
62 Ok(info) => compaction_results.push(info),
63 Err(e) => {
64 error!(
65 "Failed to compact adjacency {} <- {}: {}",
66 label, edge_type, e
67 );
68 }
69 }
70 }
71 }
72
73 metrics::counter!("uni_compaction_runs_total").increment(1);
74 metrics::histogram!("uni_compaction_duration_seconds")
75 .record(start.elapsed().as_secs_f64());
76
77 Ok(compaction_results)
78 }
79
80 #[instrument(skip(self), fields(rows_processed, duration_ms), level = "info")]
81 pub async fn compact_vertices(&self, label: &str) -> Result<()> {
82 let start = std::time::Instant::now();
83 let schema_manager = self.storage.schema_manager();
84 let schema = schema_manager.schema();
85
86 let label_props = schema
87 .properties
88 .get(label)
89 .ok_or_else(|| anyhow!("Label not found"))?;
90
91 let crdt_props: HashSet<String> = label_props
93 .iter()
94 .filter(|(_, meta)| matches!(meta.r#type, DataType::Crdt(_)))
95 .map(|(name, _)| name.clone())
96 .collect();
97
98 let dataset = self.storage.vertex_dataset(label)?;
99 let backend = self.storage.backend();
100 let table_name = dataset.table_name();
101
102 if !backend.table_exists(&table_name).await.unwrap_or(false) {
104 info!("No vertex data to compact for label '{}'", label);
105 return Ok(());
106 }
107
108 let row_count = backend.count_rows(&table_name, None).await?;
116 crate::storage::delta::check_oom_guard(
117 row_count,
118 self.storage.config.max_compaction_rows,
119 label,
120 "vertices",
121 )?;
122
123 info!(
124 label = %label,
125 row_count,
126 estimated_bytes = row_count * 200,
127 "Starting vertex compaction"
128 );
129
130 let _write_guard = backend.lock_table_for_write(&table_name).await;
136
137 use crate::backend::types::ScanRequest;
138 let batches: Vec<RecordBatch> = backend.scan(ScanRequest::all(&table_name)).await?;
139
140 let mut vertex_state: HashMap<Vid, (Properties, bool)> = HashMap::new();
142 let mut vertex_versions: HashMap<Vid, u64> = HashMap::new();
143 let mut vertex_labels: HashMap<Vid, Vec<String>> = HashMap::new();
144
145 let mut rows_processed = 0;
146
147 for batch in batches {
148 rows_processed += batch.num_rows();
149 let vid_col = batch
150 .column_by_name("_vid")
151 .unwrap()
152 .as_any()
153 .downcast_ref::<UInt64Array>()
154 .unwrap();
155 let ver_col = batch
156 .column_by_name("_version")
157 .unwrap()
158 .as_any()
159 .downcast_ref::<UInt64Array>()
160 .unwrap();
161 let del_col = batch
162 .column_by_name("_deleted")
163 .unwrap()
164 .as_any()
165 .downcast_ref::<arrow_array::BooleanArray>()
166 .unwrap();
167
168 let labels_col = batch
170 .column_by_name("_labels")
171 .and_then(|c| c.as_any().downcast_ref::<arrow_array::ListArray>());
172
173 for i in 0..batch.num_rows() {
174 let vid = Vid::from(vid_col.value(i));
175 let version = ver_col.value(i);
176 let deleted = del_col.value(i);
177
178 if let Some(list_arr) = labels_col
180 && version >= *vertex_versions.entry(vid).or_insert(0)
181 {
182 let labels = crate::storage::arrow_convert::labels_from_list_array(list_arr, i);
183 if !labels.is_empty() {
184 vertex_labels.insert(vid, labels);
185 }
186 }
187
188 let current_entry = vertex_state
189 .entry(vid)
190 .or_insert((Properties::new(), false));
191 let current_version = vertex_versions.entry(vid).or_insert(0);
192
193 if deleted {
200 if version >= *current_version {
201 current_entry.1 = true;
202 current_entry.0.clear(); *current_version = version;
204 }
205 continue;
206 }
207
208 let mut row_props = Properties::new();
211 let mut null_props = Vec::new(); for (name, meta) in label_props {
213 if let Some(col) = batch.column_by_name(name) {
214 if col.is_null(i) {
215 null_props.push(name.clone());
217 } else {
218 let val = crate::storage::value_codec::decode_column_value(
219 col.as_ref(),
220 &meta.r#type,
221 i,
222 crate::storage::value_codec::CrdtDecodeMode::Strict,
223 )?;
224 row_props.insert(name.clone(), val);
225 }
226 }
227 }
228
229 Self::merge_row_into_state(
230 row_props,
231 null_props,
232 version,
233 current_entry,
234 current_version,
235 &crdt_props,
236 self.storage.plugin_registry().map(|a| a.as_ref()),
237 )?;
238 }
239 }
240
241 let mut valid_vertices = Vec::new();
243 let mut valid_versions = Vec::new();
244 let mut valid_deleted = Vec::new(); for (vid, (props, deleted)) in vertex_state {
249 if !deleted {
250 let labels = vertex_labels.remove(&vid).unwrap_or_default();
251 valid_vertices.push((vid, labels, props));
252 valid_versions.push(vertex_versions[&vid]);
253 valid_deleted.push(false);
254 }
255 }
256
257 if !valid_vertices.is_empty() {
258 let batch = dataset.build_record_batch(
259 &valid_vertices,
260 &valid_deleted,
261 &valid_versions,
262 &schema,
263 )?;
264 dataset
265 .replace(self.storage.backend(), batch, &schema)
266 .await?;
267 }
268
269 let duration = start.elapsed();
270 let rows_reclaimed = rows_processed as u64 - valid_vertices.len() as u64;
271 metrics::counter!("uni_compaction_rows_reclaimed_total", "type" => "vertex")
272 .increment(rows_reclaimed);
273
274 tracing::Span::current().record("rows_processed", rows_processed);
275 tracing::Span::current().record("duration_ms", duration.as_millis());
276 info!(
277 rows = rows_processed,
278 duration_ms = duration.as_millis(),
279 "Vertex compaction completed"
280 );
281
282 metrics::histogram!("uni_compaction_duration_seconds", "type" => "vertex")
283 .record(duration.as_secs_f64());
284
285 Ok(())
286 }
287
288 fn merge_crdt_values(
289 a: &Value,
290 b: &Value,
291 registry: Option<&uni_plugin::PluginRegistry>,
292 ) -> Result<Value> {
293 if a.is_null() {
294 return Ok(b.clone());
295 }
296 if b.is_null() {
297 return Ok(a.clone());
298 }
299 let mut crdt_a: Crdt = serde_json::from_value(a.clone().into())?;
300 let crdt_b: Crdt = serde_json::from_value(b.clone().into())?;
301 match registry {
304 Some(reg) => crdt_a
305 .merge_via_registry(&crdt_b, reg)
306 .map_err(|e| anyhow::anyhow!("{e}"))?,
307 None => crdt_a
308 .try_merge(&crdt_b)
309 .map_err(|e| anyhow::anyhow!("{e}"))?,
310 }
311 Ok(Value::from(serde_json::to_value(crdt_a)?))
312 }
313
314 fn merge_row_into_state(
316 row_props: Properties,
317 null_props: Vec<String>,
318 version: u64,
319 current_entry: &mut (Properties, bool),
320 current_version: &mut u64,
321 crdt_props: &HashSet<String>,
322 registry: Option<&uni_plugin::PluginRegistry>,
323 ) -> Result<()> {
324 if version > *current_version {
325 *current_version = version;
327 current_entry.1 = false;
328
329 for (k, v) in row_props {
330 if crdt_props.contains(&k) {
331 let existing = current_entry.0.entry(k.clone()).or_insert(Value::Null);
332 *existing = Self::merge_crdt_values(existing, &v, registry)?;
333 } else {
334 current_entry.0.insert(k, v);
335 }
336 }
337
338 for null_prop in &null_props {
340 if !crdt_props.contains(null_prop) {
341 current_entry.0.remove(null_prop);
342 }
343 }
344 } else if version == *current_version {
345 current_entry.1 = false;
347 for (k, v) in row_props {
348 if crdt_props.contains(&k) {
349 let existing = current_entry.0.entry(k.clone()).or_insert(Value::Null);
350 *existing = Self::merge_crdt_values(existing, &v, registry)?;
351 } else {
352 current_entry.0.insert(k, v);
353 }
354 }
355 } else {
356 if !current_entry.1 {
358 for (k, v) in row_props {
359 if crdt_props.contains(&k) {
360 let existing = current_entry.0.entry(k.clone()).or_insert(Value::Null);
361 *existing = Self::merge_crdt_values(existing, &v, registry)?;
362 }
363 }
364 }
365 }
366 Ok(())
367 }
368
369 #[instrument(skip(self), fields(delta_count, duration_ms), level = "info")]
370 pub async fn compact_adjacency(
371 &self,
372 edge_type: &str,
373 label: &str,
374 direction: &str,
375 ) -> Result<CompactionInfo> {
376 let start = std::time::Instant::now();
377 let schema = self.storage.schema_manager().schema();
378
379 let delta_ds = self.storage.delta_dataset(edge_type, direction)?;
381 let deltas = delta_ds
382 .scan_all_backend(self.storage.backend(), &schema)
383 .await?;
384
385 let delta_count = deltas.len();
386 tracing::Span::current().record("delta_count", delta_count);
387
388 if deltas.is_empty() {
389 return Ok(CompactionInfo {
391 edge_type: edge_type.to_string(),
392 direction: direction.to_string(),
393 });
394 }
395
396 let mut delta_map: HashMap<Vid, Vec<L1Entry>> = HashMap::new();
400 for entry in &deltas {
401 let key = if direction == "fwd" {
402 entry.src_vid
403 } else {
404 entry.dst_vid
405 };
406 delta_map.entry(key).or_default().push(entry.clone());
407 }
408
409 for ops in delta_map.values_mut() {
412 ops.sort_by_key(|e| e.version);
413 }
414
415 let adj_ds = self
417 .storage
418 .adjacency_dataset(edge_type, label, direction)?;
419
420 let mut src_vid_builder = UInt64Builder::new();
431 let mut neighbors_builder = ListBuilder::new(UInt64Builder::new());
432 let mut edge_ids_builder = ListBuilder::new(UInt64Builder::new());
433
434 let mut processed_vids = HashSet::new();
435
436 let backend = self.storage.backend();
438 let adj_table_name = adj_ds.table_name();
439 if backend.table_exists(&adj_table_name).await.unwrap_or(false) {
440 let adj_row_count = backend.count_rows(&adj_table_name, None).await?;
441 crate::storage::delta::check_oom_guard(
442 adj_row_count,
443 self.storage.config.max_compaction_rows,
444 &format!("{}_{}", edge_type, label),
445 direction,
446 )?;
447
448 info!(
449 edge_type = %edge_type,
450 label = %label,
451 direction = %direction,
452 adj_row_count,
453 delta_count,
454 estimated_bytes = adj_row_count * 100 + delta_count * ENTRY_SIZE_ESTIMATE,
455 "Starting adjacency compaction"
456 );
457
458 use crate::backend::types::ScanRequest;
459 let batches: Vec<RecordBatch> = backend.scan(ScanRequest::all(&adj_table_name)).await?;
460
461 for batch in batches {
462 let src_col = batch
463 .column_by_name("src_vid")
464 .ok_or(anyhow!("Missing src_vid"))?
465 .as_any()
466 .downcast_ref::<UInt64Array>()
467 .ok_or(anyhow!("Invalid src_vid"))?;
468 let neighbors_col = batch
469 .column_by_name("neighbors")
470 .ok_or(anyhow!("Missing neighbors"))?
471 .as_any()
472 .downcast_ref::<ListArray>()
473 .ok_or(anyhow!("Invalid neighbors"))?;
474 let edge_ids_col = batch
475 .column_by_name("edge_ids")
476 .ok_or(anyhow!("Missing edge_ids"))?
477 .as_any()
478 .downcast_ref::<ListArray>()
479 .ok_or(anyhow!("Invalid edge_ids"))?;
480
481 for i in 0..batch.num_rows() {
482 let vid = Vid::from(src_col.value(i));
483 processed_vids.insert(vid);
484
485 let n_list = neighbors_col.value(i);
487 let n_array = n_list.as_any().downcast_ref::<UInt64Array>().unwrap();
488 let e_list = edge_ids_col.value(i);
489 let e_array = e_list.as_any().downcast_ref::<UInt64Array>().unwrap();
490
491 let mut current_edges: HashMap<Eid, Vid> = HashMap::new();
492 for j in 0..n_array.len() {
493 current_edges
494 .insert(Eid::from(e_array.value(j)), Vid::from(n_array.value(j)));
495 }
496
497 if let Some(ops) = delta_map.get(&vid) {
498 apply_deltas_to_edges(&mut current_edges, ops, direction);
499 }
500
501 append_edges_to_builders(
502 vid,
503 ¤t_edges,
504 &mut src_vid_builder,
505 &mut neighbors_builder,
506 &mut edge_ids_builder,
507 );
508 }
509 }
510 }
511
512 for (vid, ops) in delta_map {
514 if processed_vids.contains(&vid) {
515 continue;
516 }
517
518 let mut current_edges: HashMap<Eid, Vid> = HashMap::new();
519 apply_deltas_to_edges(&mut current_edges, &ops, direction);
520
521 append_edges_to_builders(
522 vid,
523 ¤t_edges,
524 &mut src_vid_builder,
525 &mut neighbors_builder,
526 &mut edge_ids_builder,
527 );
528 }
529
530 {
537 let src_arr = Arc::new(src_vid_builder.finish());
538 let neighbors_arr = Arc::new(neighbors_builder.finish());
539 let edge_ids_arr = Arc::new(edge_ids_builder.finish());
540
541 let schema = adj_ds.get_arrow_schema();
542 let batch = RecordBatch::try_new(schema, vec![src_arr, neighbors_arr, edge_ids_arr])?;
543
544 adj_ds.replace(self.storage.backend(), batch).await?;
546 }
547
548 if !deltas.is_empty() {
553 info!(
554 "Clearing Delta L1 for edge_type={} direction={} after compaction (incorporated {} ops)",
555 edge_type,
556 direction,
557 deltas.len()
558 );
559
560 #[cfg(debug_assertions)]
565 {
566 use crate::storage::main_edge::MainEdgeDataset;
567
568 let delta_eids: std::collections::HashSet<Eid> =
569 deltas.iter().map(|e| e.eid).collect();
570
571 for eid in delta_eids {
572 let main_edge_exists =
573 MainEdgeDataset::exists_by_eid(self.storage.backend(), eid)
574 .await
575 .unwrap_or(false);
576
577 debug_assert!(
578 main_edge_exists,
579 "EID {} from Delta L1 not found in main_edges after compaction. \
580 This indicates edge properties were not dual-written during flush.",
581 eid.as_u64()
582 );
583 }
584 }
585
586 let clear_hwm = deltas.iter().map(|e| e.version).max().unwrap_or(0);
596 let delta_ds = self.storage.delta_dataset(edge_type, direction)?;
597 delta_ds
598 .delete_up_to_version(self.storage.backend(), clear_hwm)
599 .await?;
600 }
601
602 let duration = start.elapsed();
603 tracing::Span::current().record("duration_ms", duration.as_millis());
604 info!(
605 delta_count,
606 duration_ms = duration.as_millis(),
607 "Adjacency compaction completed"
608 );
609
610 metrics::histogram!("uni_compaction_duration_seconds", "type" => "adjacency")
611 .record(duration.as_secs_f64());
612
613 Ok(CompactionInfo {
614 edge_type: edge_type.to_string(),
615 direction: direction.to_string(),
616 })
617 }
618}
619
620fn apply_deltas_to_edges(current_edges: &mut HashMap<Eid, Vid>, ops: &[L1Entry], direction: &str) {
622 for op in ops {
623 match op.op {
624 Op::Insert => {
625 let neighbor = if direction == "fwd" {
626 op.dst_vid
627 } else {
628 op.src_vid
629 };
630 current_edges.insert(op.eid, neighbor);
631 }
632 Op::Delete => {
633 current_edges.remove(&op.eid);
634 }
635 }
636 }
637}
638
639fn append_edges_to_builders(
641 vid: Vid,
642 current_edges: &HashMap<Eid, Vid>,
643 src_vid_builder: &mut UInt64Builder,
644 neighbors_builder: &mut ListBuilder<UInt64Builder>,
645 edge_ids_builder: &mut ListBuilder<UInt64Builder>,
646) {
647 if current_edges.is_empty() {
648 return;
649 }
650 src_vid_builder.append_value(vid.as_u64());
651
652 let mut sorted_eids: Vec<_> = current_edges.keys().cloned().collect();
653 sorted_eids.sort();
654
655 for eid in sorted_eids {
656 let neighbor = current_edges[&eid];
657 neighbors_builder.values().append_value(neighbor.as_u64());
658 edge_ids_builder.values().append_value(eid.as_u64());
659 }
660 neighbors_builder.append(true);
661 edge_ids_builder.append(true);
662}
663
664#[derive(Debug, Clone)]
667pub struct CompactionInfo {
668 pub edge_type: String,
669 pub direction: String,
670}
671
672#[cfg(test)]
673mod ws_d_crdt_tests {
674 use std::sync::Arc;
683 use std::sync::atomic::{AtomicUsize, Ordering};
684
685 use uni_common::Value;
686 use uni_crdt::{Crdt, GCounter};
687 use uni_plugin::traits::crdt::{CrdtKind, CrdtKindProvider, CrdtOp, CrdtState, ScalarValue};
688 use uni_plugin::{
689 Capability, CapabilitySet, FnError, PluginId, PluginRegistrar, PluginRegistry,
690 };
691
692 #[derive(Default)]
693 struct CountingProvider {
694 calls: AtomicUsize,
695 }
696 impl CrdtKindProvider for CountingProvider {
697 fn kind(&self) -> CrdtKind {
698 CrdtKind::new("uni-crdt:g-counter")
699 }
700 fn empty(&self) -> Box<dyn CrdtState> {
701 Box::new(St {
702 inner: Crdt::GCounter(GCounter::new()),
703 })
704 }
705 fn from_persisted(&self, bytes: &[u8]) -> Result<Box<dyn CrdtState>, FnError> {
706 self.calls.fetch_add(1, Ordering::SeqCst);
707 let inner =
708 Crdt::from_msgpack(bytes).map_err(|e| FnError::new(0xA01, format!("{e}")))?;
709 Ok(Box::new(St { inner }))
710 }
711 }
712 struct St {
713 inner: Crdt,
714 }
715 impl CrdtState for St {
716 fn as_any(&self) -> &dyn std::any::Any {
717 self
718 }
719 fn apply(&mut self, _op: &CrdtOp) -> Result<(), FnError> {
720 Ok(())
721 }
722 fn merge(&mut self, other: &dyn CrdtState) -> Result<(), FnError> {
723 let o = other
724 .as_any()
725 .downcast_ref::<St>()
726 .ok_or_else(|| FnError::new(0xA10, "type mismatch"))?;
727 self.inner
728 .try_merge(&o.inner)
729 .map_err(|e| FnError::new(0xA11, format!("{e}")))
730 }
731 fn value(&self) -> Result<ScalarValue, FnError> {
732 Ok(ScalarValue::Utf8(Some(self.inner.type_name().to_owned())))
733 }
734 fn persist(&self) -> Result<Vec<u8>, FnError> {
735 self.inner
736 .to_msgpack()
737 .map_err(|e| FnError::new(0xA12, format!("{e}")))
738 }
739 }
740
741 fn gcounter(replica: &str, by: u64) -> Value {
742 let mut g = GCounter::new();
743 g.increment(replica, by);
744 Value::from(serde_json::to_value(Crdt::GCounter(g)).unwrap())
745 }
746
747 #[test]
748 fn compaction_merge_routes_through_registry() {
749 let registry = PluginRegistry::new();
750 let provider = Arc::new(CountingProvider::default());
751 let caps = CapabilitySet::from_iter_of([Capability::Crdt]);
752 let mut r = PluginRegistrar::new(PluginId::new("test.counting"), &caps, ®istry);
753 r.crdt_kind(
754 CrdtKind::new("uni-crdt:g-counter"),
755 Arc::clone(&provider) as Arc<dyn CrdtKindProvider>,
756 )
757 .unwrap();
758 r.commit_to_registry().unwrap();
759
760 let a = gcounter("r1", 5); let b = gcounter("r2", 7); let merged = super::Compactor::merge_crdt_values(&a, &b, Some(®istry)).unwrap();
763
764 assert!(
765 provider.calls.load(Ordering::SeqCst) > 0,
766 "compaction merge must route through the registered provider"
767 );
768 let crdt: Crdt = serde_json::from_value(merged.into()).unwrap();
769 match crdt {
770 Crdt::GCounter(g) => assert_eq!(g.value(), 12, "5 + 7 = 12"),
771 other => panic!("expected GCounter, got {other:?}"),
772 }
773 }
774
775 #[test]
776 fn compaction_merge_falls_back_to_native_without_registry() {
777 let a = gcounter("r1", 3);
778 let b = gcounter("r2", 4);
779 let merged = super::Compactor::merge_crdt_values(&a, &b, None).unwrap();
780 let crdt: Crdt = serde_json::from_value(merged.into()).unwrap();
781 match crdt {
782 Crdt::GCounter(g) => assert_eq!(g.value(), 7, "3 + 4 = 7 native fallback"),
783 other => panic!("expected GCounter, got {other:?}"),
784 }
785 }
786}