1use std::collections::{BTreeMap, HashMap, HashSet};
16
17use crate::compaction::CompactionOutput;
18use crate::compaction::picker::PickerOutput;
19use crate::sst::file::FileHandle;
20
21#[derive(Debug, Clone)]
26pub(crate) struct CompactionUnit {
27 pub(crate) inputs: Vec<FileHandle>,
29 pub(crate) outputs: Vec<CompactionOutput>,
31 pub(crate) time_window_size: i64,
33 pub(crate) max_file_size: Option<usize>,
35}
36
37impl CompactionUnit {
38 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 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 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 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
134fn 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 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}