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 memory once per input, even when several outputs read it.
113    pub(crate) fn estimated_memory_bytes(&self) -> u64 {
114        if self.outputs.is_empty() {
115            return 0;
116        }
117        self.inputs.iter().fold(0u64, |total, file| {
118            total.saturating_add(file.meta_ref().max_row_group_uncompressed_size)
119        })
120    }
121}
122
123/// Finds a dependency group's representative while shortening the parent chain.
124fn root(parents: &mut [usize], mut index: usize) -> usize {
125    while parents[index] != index {
126        parents[index] = parents[parents[index]];
127        index = parents[index];
128    }
129    index
130}
131
132#[cfg(test)]
133mod tests {
134    use store_api::storage::FileId;
135
136    use super::*;
137    use crate::compaction::test_util::new_file_handle;
138
139    /// Builds an L1 output that preserves deletion markers for grouping tests.
140    fn output(inputs: Vec<FileHandle>) -> CompactionOutput {
141        CompactionOutput {
142            inputs,
143            output_level: 1,
144            filter_deleted: false,
145            output_time_range: None,
146        }
147    }
148
149    #[test]
150    fn test_compaction_unit_transitive_inputs() {
151        let files: Vec<_> = (0..3)
152            .map(|_| new_file_handle(FileId::random(), 0, 100, 0))
153            .collect();
154        let units = CompactionUnit::from_picker(PickerOutput {
155            outputs: vec![
156                output(vec![files[0].clone()]),
157                output(vec![files[2].clone()]),
158                output(vec![files[1].clone()]),
159                output(vec![files[0].clone(), files[1].clone()]),
160            ],
161            time_window_size: 100,
162            max_file_size: Some(1024),
163            ..Default::default()
164        })
165        .unwrap();
166        assert_eq!(2, units.len());
167        assert_eq!(3, units[0].outputs.len());
168        assert_eq!(2, units[0].inputs.len());
169        assert_eq!(files[2].file_id(), units[1].inputs[0].file_id());
170        assert_eq!(Some(1024), units[0].max_file_size);
171        assert_eq!(100, units[1].time_window_size);
172    }
173
174    #[test]
175    fn test_compaction_unit_group_inherits_highest_output_priority() {
176        let files: Vec<_> = (0..3)
177            .map(|_| new_file_handle(FileId::random(), 0, 100, 0))
178            .collect();
179        let units = CompactionUnit::from_picker(PickerOutput {
180            outputs: [0, 1, 1, 2, 0]
181                .into_iter()
182                .map(|index| output(vec![files[index].clone()]))
183                .collect(),
184            ..Default::default()
185        })
186        .unwrap();
187        assert_eq!(
188            vec![files[0].file_id(), files[2].file_id(), files[1].file_id()],
189            units
190                .iter()
191                .map(|unit| unit.inputs[0].file_id())
192                .collect::<Vec<_>>()
193        );
194        assert_eq!(
195            vec![2, 1, 2],
196            units
197                .iter()
198                .map(|unit| unit.outputs.len())
199                .collect::<Vec<_>>()
200        );
201    }
202
203    #[test]
204    fn test_compaction_unit_expiration_precedes_merge_units() {
205        let files: Vec<_> = (0..3)
206            .map(|_| new_file_handle(FileId::random(), 0, 100, 0))
207            .collect();
208        let units = CompactionUnit::from_picker(PickerOutput {
209            outputs: vec![
210                output(vec![files[0].clone()]),
211                output(vec![files[1].clone()]),
212            ],
213            expired_ssts: vec![files[2].clone()],
214            ..Default::default()
215        })
216        .unwrap();
217
218        assert_eq!(3, units.len());
219        assert!(units[0].outputs.is_empty());
220        assert_eq!(files[2].file_id(), units[0].inputs[0].file_id());
221        assert_eq!(files[1].file_id(), units[1].inputs[0].file_id());
222        assert_eq!(files[0].file_id(), units[2].inputs[0].file_id());
223    }
224
225    #[test]
226    fn test_compaction_unit_expiration_inputs() {
227        let file = new_file_handle(FileId::random(), 0, 100, 0);
228        assert!(
229            CompactionUnit::from_picker(PickerOutput {
230                outputs: vec![output(vec![file.clone()])],
231                expired_ssts: vec![file.clone()],
232                ..Default::default()
233            })
234            .is_none()
235        );
236        let units = CompactionUnit::from_picker(PickerOutput {
237            expired_ssts: vec![file.clone(), file],
238            ..Default::default()
239        })
240        .unwrap();
241        assert_eq!(1, units.len());
242        assert_eq!(1, units[0].inputs.len());
243        assert!(units[0].outputs.is_empty());
244        assert_eq!(0, units[0].estimated_memory_bytes());
245        assert!(
246            CompactionUnit::from_picker(PickerOutput::default())
247                .unwrap()
248                .is_empty()
249        );
250    }
251}
mito2/compaction/unit.rs

mito2/compaction/
unit.rs
