erigon_seg/writer/
merge.rs1use std::cmp::Reverse;
15use std::collections::BinaryHeap;
16use std::path::Path;
17
18use super::domain::{DomainOptions, DomainPaths, DomainWriter};
19use crate::error::{Error, Result};
20use crate::seg::Seg;
21
22#[derive(Debug, Clone, Copy)]
24pub struct MergeOptions {
25 pub domain: DomainOptions,
27 pub drop_deleted: bool,
30 pub range_from: Option<u64>,
34}
35
36impl Default for MergeOptions {
37 fn default() -> MergeOptions {
38 MergeOptions {
39 domain: DomainOptions::default(),
40 drop_deleted: true,
41 range_from: None,
42 }
43 }
44}
45
46pub fn merge(
53 inputs: &[impl AsRef<Path>],
54 out_kv: impl AsRef<Path>,
55 opts: MergeOptions,
56) -> Result<DomainPaths> {
57 let out_kv = out_kv.as_ref();
58 if inputs.is_empty() {
59 return Err(Error::format("merge: no input files"));
60 }
61
62 let mut order: Vec<usize> = (0..inputs.len()).collect();
64 order.sort_by_key(|&i| (parse_from(inputs[i].as_ref()).unwrap_or(0), i));
65
66 let segs: Vec<Seg> = order
67 .iter()
68 .map(|&i| Seg::open(inputs[i].as_ref()))
69 .collect::<Result<Vec<_>>>()?;
70 let mut getters: Vec<_> = segs.iter().map(|s| s.getter()).collect();
71
72 let mut heads: Vec<Option<(Vec<u8>, Vec<u8>)>> = vec![None; segs.len()];
74 let mut heap: BinaryHeap<Reverse<(Vec<u8>, usize)>> = BinaryHeap::new();
75 for (i, g) in getters.iter_mut().enumerate() {
76 if g.has_next() {
77 let k = g.next();
78 let v = if g.has_next() { g.next() } else { Vec::new() };
79 heap.push(Reverse((k.clone(), i)));
80 heads[i] = Some((k, v));
81 }
82 }
83
84 let range_from = opts.range_from.or_else(|| parse_from(out_kv));
85 let drop_at_zero = opts.drop_deleted && range_from == Some(0);
86
87 let mut writer = DomainWriter::create(out_kv, opts.domain)?;
88 while let Some(Reverse((min_key, idx0))) = heap.pop() {
89 let mut best_idx = idx0;
92 let mut best_val = heads[idx0].take().expect("head present for heap entry").1;
93 let mut advance: Vec<usize> = vec![idx0];
94 while let Some(Reverse((k, _))) = heap.peek() {
95 if *k != min_key {
96 break;
97 }
98 let Reverse((_, idx)) = heap.pop().unwrap();
99 let val = heads[idx].take().expect("head present for heap entry").1;
100 if idx > best_idx {
101 best_idx = idx;
102 best_val = val;
103 }
104 advance.push(idx);
105 }
106
107 for idx in advance {
109 if getters[idx].has_next() {
110 let k = getters[idx].next();
111 let v = if getters[idx].has_next() {
112 getters[idx].next()
113 } else {
114 Vec::new()
115 };
116 heap.push(Reverse((k.clone(), idx)));
117 heads[idx] = Some((k, v));
118 }
119 }
120
121 if drop_at_zero && best_val.is_empty() {
122 continue; }
124 writer.add(&min_key, &best_val)?;
125 }
126
127 writer.finish()
128}
129
130fn parse_from(path: &Path) -> Option<u64> {
132 let name = path.file_name()?.to_string_lossy();
133 for seg in name.split('.') {
134 if let Some((a, b)) = seg.split_once('-')
135 && let (Ok(a), Ok(_)) = (a.parse::<u64>(), b.parse::<u64>())
136 {
137 return Some(a);
138 }
139 }
140 None
141}