Skip to main content

erigon_seg/writer/
merge.rs

1//! K-way merge of several domain `.kv` files into one.
2//!
3//! Inputs are sorted, unique-keyed domain files. The merge emits each distinct key once,
4//! in order, taking the value from the **newest** input that contains it (newest wins,
5//! matching erigon's step-range override semantics and this crate's query-time
6//! newest-wins in a multi-file reader).
7//!
8//! **Deleted entries are dropped** following erigon exactly: a key is omitted iff the
9//! merged range starts at step 0 *and* the winning value is empty (`r.values.from == 0
10//! && len(val) == 0`). Outside a from-zero merge an empty value is a legitimate stored
11//! value (e.g. an empty account) and is preserved. The range start is taken from
12//! [`MergeOptions::range_from`], or parsed from the output filename's `<from>-<to>`.
13
14use 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/// Options for [`merge`].
23#[derive(Debug, Clone, Copy)]
24pub struct MergeOptions {
25    /// `.bt`/`.kvei` options for the merged output.
26    pub domain: DomainOptions,
27    /// Honor erigon's "empty value means deletion" rule (only triggers when the merged
28    /// range starts at step 0). Default `true`.
29    pub drop_deleted: bool,
30    /// The merged range's start step. If `None`, parsed from the output `.kv` filename's
31    /// `<from>-<to>` segment; if that fails too, deletion-dropping is disabled (no key is
32    /// ever dropped), which is the safe choice.
33    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
46/// Merge `inputs` (domain `.kv` paths) into `out_kv`, also building the sibling `.bt`
47/// and—if `opts.domain.salt` is set—`.kvei`.
48///
49/// Inputs are reordered oldest→newest by the `<from>` step parsed from their filenames
50/// (stable, so files without a parseable step keep their given order); the last is
51/// treated as newest. Returns the paths written.
52pub 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    // Order oldest -> newest by parsed `from` (stable for unparseable names).
63    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    // Per-input current head (key, value); the heap orders inputs by head key.
73    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        // Winner so far is the entry just popped; scan all other inputs with this key and
90        // keep the one from the newest input (highest index).
91        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        // Refill the advanced inputs.
108        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; // deletion in a from-zero merge
123        }
124        writer.add(&min_key, &best_val)?;
125    }
126
127    writer.finish()
128}
129
130/// Parse the `<from>` step from a `…<from>-<to>.<ext>` filename, as a number of steps.
131fn 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}