Skip to main content

mito2/compaction/
unit.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, HashMap, HashSet};
16
17use crate::compaction::CompactionOutput;
18use crate::compaction::picker::PickerOutput;
19use crate::sst::file::FileHandle;
20
21/// A complete input dependency group, merged and published atomically.
22///
23/// Logical outputs retain their original input sets and time bounds. In particular,
24/// splitting physical SSTs must never split a strict-window input dependency group.
25#[derive(Debug, Clone)]
26pub(crate) struct CompactionUnit {
27    /// Distinct input SSTs reserved and removed together.
28    pub(crate) inputs: Vec<FileHandle>,
29    /// Original logical outputs; empty for expiration-only maintenance.
30    pub(crate) outputs: Vec<CompactionOutput>,
31    /// Compaction window size in seconds.
32    pub(crate) time_window_size: i64,
33    /// Maximum physical output SST size in bytes.
34    pub(crate) max_file_size: Option<usize>,
35}
36
37impl CompactionUnit {
38    /// Groups outputs by the transitive closure of shared input FileIds.
39    /// Returns `None` for an inconsistent plan that both expires and merges a file.
40    /// Expiration maintenance runs before merge units, matching the legacy pre-merge removal.
41    /// Merge units follow the legacy compactor's tail-first execution order; a shared-input
42    /// group inherits the position of its highest-priority output.
43    pub(crate) fn from_picker(output: PickerOutput) -> Option<Vec<Self>> {
44        let PickerOutput {
45            outputs,
46            expired_ssts,
47            time_window_size,
48            max_file_size,
49        } = output;
50        let mut parents: Vec<_> = (0..outputs.len()).collect();
51        let mut owners = HashMap::new();
52        for (index, output) in outputs.iter().enumerate() {
53            for file in &output.inputs {
54                if let Some(&previous) = owners.get(&file.file_id()) {
55                    let a = root(&mut parents, previous);
56                    let b = root(&mut parents, index);
57                    // Higher output indices are consumed first by the legacy compactor.
58                    parents[a.min(b)] = a.max(b);
59                } else {
60                    owners.insert(file.file_id(), index);
61                }
62            }
63        }
64        if expired_ssts
65            .iter()
66            .any(|f| owners.contains_key(&f.file_id()))
67        {
68            return None;
69        }
70
71        let mut groups: BTreeMap<usize, Vec<CompactionOutput>> = BTreeMap::new();
72        for (index, output) in outputs.into_iter().enumerate() {
73            groups
74                .entry(root(&mut parents, index))
75                .or_default()
76                .push(output);
77        }
78        let mut units = Vec::with_capacity(groups.len() + 1);
79        if !expired_ssts.is_empty() {
80            let mut seen = HashSet::new();
81            units.push(Self {
82                inputs: expired_ssts
83                    .into_iter()
84                    .filter(|f| seen.insert(f.file_id()))
85                    .collect(),
86                outputs: Vec::new(),
87                time_window_size,
88                max_file_size,
89            });
90        }
91        for outputs in groups.into_values().rev() {
92            let mut seen = HashSet::new();
93            let inputs: Vec<_> = outputs
94                .iter()
95                .flat_map(|o| &o.inputs)
96                .filter(|f| seen.insert(f.file_id()))
97                .cloned()
98                .collect();
99            if inputs.is_empty() {
100                continue;
101            }
102            units.push(Self {
103                inputs,
104                outputs,
105                time_window_size,
106                max_file_size,
107            });
108        }
109        Some(units)
110    }
111
112    /// Estimates merge work from distinct input bytes. Expiration-only units do
113    /// not read SSTs; admission assigns them its minimum positive service charge.
114    pub(crate) fn estimated_input_bytes(&self) -> u64 {
115        if self.outputs.is_empty() {
116            return 0;
117        }
118        self.inputs
119            .iter()
120            .fold(0u64, |total, file| total.saturating_add(file.size()))
121    }
122
123    /// Estimates merge memory once per input, even when several outputs read it.
124    pub(crate) fn estimated_memory_bytes(&self) -> u64 {
125        if self.outputs.is_empty() {
126            return 0;
127        }
128        self.inputs.iter().fold(0u64, |total, file| {
129            total.saturating_add(file.meta_ref().max_row_group_uncompressed_size)
130        })
131    }
132}
133
134/// Finds a dependency group's representative while shortening the parent chain.
135fn root(parents: &mut [usize], mut index: usize) -> usize {
136    while parents[index] != index {
137        parents[index] = parents[parents[index]];
138        index = parents[index];
139    }
140    index
141}
142
143#[cfg(test)]
144mod tests {
145    use store_api::storage::FileId;
146
147    use super::*;
148    use crate::compaction::test_util::{new_file_handle, new_file_handle_with_size_and_sequence};
149
150    /// Builds an L1 output that preserves deletion markers for grouping tests.
151    fn output(inputs: Vec<FileHandle>) -> CompactionOutput {
152        CompactionOutput {
153            inputs,
154            output_level: 1,
155            filter_deleted: false,
156            output_time_range: None,
157        }
158    }
159
160    #[test]
161    fn test_compaction_unit_transitive_inputs() {
162        let files: Vec<_> = (0..3)
163            .map(|i| {
164                new_file_handle_with_size_and_sequence(FileId::random(), 0, 100, 0, 0, (i + 1) * 10)
165            })
166            .collect();
167        let units = CompactionUnit::from_picker(PickerOutput {
168            outputs: vec![
169                output(vec![files[0].clone()]),
170                output(vec![files[2].clone()]),
171                output(vec![files[1].clone()]),
172                output(vec![files[0].clone(), files[1].clone()]),
173            ],
174            time_window_size: 100,
175            max_file_size: Some(1024),
176            ..Default::default()
177        })
178        .unwrap();
179        assert_eq!(2, units.len());
180        assert_eq!(3, units[0].outputs.len());
181        assert_eq!(2, units[0].inputs.len());
182        assert_eq!(files[2].file_id(), units[1].inputs[0].file_id());
183        assert_eq!(30, units[0].estimated_input_bytes());
184        assert_eq!(30, units[1].estimated_input_bytes());
185        assert_eq!(Some(1024), units[0].max_file_size);
186        assert_eq!(100, units[1].time_window_size);
187    }
188
189    #[test]
190    fn test_compaction_unit_group_inherits_highest_output_priority() {
191        let files: Vec<_> = (0..3)
192            .map(|_| new_file_handle(FileId::random(), 0, 100, 0))
193            .collect();
194        let units = CompactionUnit::from_picker(PickerOutput {
195            outputs: [0, 1, 1, 2, 0]
196                .into_iter()
197                .map(|index| output(vec![files[index].clone()]))
198                .collect(),
199            ..Default::default()
200        })
201        .unwrap();
202        assert_eq!(
203            vec![files[0].file_id(), files[2].file_id(), files[1].file_id()],
204            units
205                .iter()
206                .map(|unit| unit.inputs[0].file_id())
207                .collect::<Vec<_>>()
208        );
209        assert_eq!(
210            vec![2, 1, 2],
211            units
212                .iter()
213                .map(|unit| unit.outputs.len())
214                .collect::<Vec<_>>()
215        );
216    }
217
218    #[test]
219    fn test_compaction_unit_expiration_precedes_merge_units() {
220        let files: Vec<_> = (0..3)
221            .map(|_| new_file_handle(FileId::random(), 0, 100, 0))
222            .collect();
223        let units = CompactionUnit::from_picker(PickerOutput {
224            outputs: vec![
225                output(vec![files[0].clone()]),
226                output(vec![files[1].clone()]),
227            ],
228            expired_ssts: vec![files[2].clone()],
229            ..Default::default()
230        })
231        .unwrap();
232
233        assert_eq!(3, units.len());
234        assert!(units[0].outputs.is_empty());
235        assert_eq!(files[2].file_id(), units[0].inputs[0].file_id());
236        assert_eq!(files[1].file_id(), units[1].inputs[0].file_id());
237        assert_eq!(files[0].file_id(), units[2].inputs[0].file_id());
238    }
239
240    #[test]
241    fn test_compaction_unit_expiration_inputs() {
242        let file = new_file_handle(FileId::random(), 0, 100, 0);
243        assert!(
244            CompactionUnit::from_picker(PickerOutput {
245                outputs: vec![output(vec![file.clone()])],
246                expired_ssts: vec![file.clone()],
247                ..Default::default()
248            })
249            .is_none()
250        );
251        let units = CompactionUnit::from_picker(PickerOutput {
252            expired_ssts: vec![file.clone(), file],
253            ..Default::default()
254        })
255        .unwrap();
256        assert_eq!(1, units.len());
257        assert_eq!(1, units[0].inputs.len());
258        assert!(units[0].outputs.is_empty());
259        assert_eq!(0, units[0].estimated_memory_bytes());
260        assert!(
261            CompactionUnit::from_picker(PickerOutput::default())
262                .unwrap()
263                .is_empty()
264        );
265    }
266}