Skip to main content

mito2/compaction/
compactor.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::num::NonZero;
16use std::sync::Arc;
17use std::time::Duration;
18
19use api::v1::region::compact_request;
20use common_meta::key::SchemaMetadataManagerRef;
21use common_telemetry::{info, warn};
22use common_time::TimeToLive;
23use either::Either;
24use itertools::Itertools;
25use object_store::manager::ObjectStoreManagerRef;
26use partition::expr::PartitionExpr;
27use serde::{Deserialize, Serialize};
28use snafu::{OptionExt, ResultExt};
29use store_api::metadata::RegionMetadataRef;
30use store_api::region_request::PathType;
31use store_api::storage::RegionId;
32
33use crate::access_layer::{
34    AccessLayer, AccessLayerRef, Metrics, OperationType, SstWriteRequest, WriteType,
35};
36use crate::cache::{CacheManager, CacheManagerRef};
37use crate::compaction::picker::{PickerOutput, new_picker};
38use crate::compaction::{CompactionOutput, CompactionSstReaderBuilder, find_dynamic_options};
39use crate::config::MitoConfig;
40use crate::error::{
41    EmptyRegionDirSnafu, InvalidPartitionExprSnafu, JoinSnafu, ObjectStoreNotFoundSnafu, Result,
42};
43use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
44use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
45use crate::metrics;
46use crate::read::FlatSource;
47use crate::region::options::RegionOptions;
48use crate::region::version::VersionRef;
49use crate::region::{ManifestContext, RegionLeaderState, RegionRoleState};
50use crate::schedule::scheduler::LocalScheduler;
51use crate::sst::FormatType;
52use crate::sst::file::FileMeta;
53use crate::sst::file_purger::LocalFilePurger;
54use crate::sst::index::intermediate::IntermediateManager;
55use crate::sst::index::puffin_manager::PuffinManagerFactory;
56use crate::sst::location::region_dir_from_table_dir;
57use crate::sst::parquet::WriteOptions;
58use crate::sst::version::{SstVersion, SstVersionRef};
59
60/// Region version for compaction that does not hold memtables.
61#[derive(Clone)]
62pub struct CompactionVersion {
63    /// Metadata of the region.
64    ///
65    /// Altering metadata isn't frequent, storing metadata in Arc to allow sharing
66    /// metadata and reuse metadata when creating a new `Version`.
67    pub(crate) metadata: RegionMetadataRef,
68    /// Options of the region.
69    pub(crate) options: RegionOptions,
70    /// SSTs of the region.
71    pub(crate) ssts: SstVersionRef,
72    /// Inferred compaction time window.
73    pub(crate) compaction_time_window: Option<Duration>,
74}
75
76impl From<VersionRef> for CompactionVersion {
77    fn from(value: VersionRef) -> Self {
78        Self {
79            metadata: value.metadata.clone(),
80            options: value.options.clone(),
81            ssts: value.ssts.clone(),
82            compaction_time_window: value.compaction_time_window,
83        }
84    }
85}
86
87/// CompactionRegion represents a region that needs to be compacted.
88/// It's the subset of MitoRegion.
89#[derive(Clone)]
90pub struct CompactionRegion {
91    pub region_id: RegionId,
92    pub region_options: RegionOptions,
93
94    pub(crate) engine_config: Arc<MitoConfig>,
95    pub(crate) region_metadata: RegionMetadataRef,
96    pub(crate) cache_manager: CacheManagerRef,
97    /// Access layer to get the table path and path type.
98    pub access_layer: AccessLayerRef,
99    pub(crate) manifest_ctx: Arc<ManifestContext>,
100    pub(crate) current_version: CompactionVersion,
101    pub(crate) file_purger: Option<Arc<LocalFilePurger>>,
102    pub(crate) ttl: Option<TimeToLive>,
103
104    /// Controls the parallelism of this compaction task. Default is 1.
105    ///
106    /// The parallel is inside this compaction task, not across different compaction tasks.
107    /// It can be different windows of the same compaction task or something like this.
108    pub max_parallelism: usize,
109}
110
111/// OpenCompactionRegionRequest represents the request to open a compaction region.
112#[derive(Debug, Clone)]
113pub struct OpenCompactionRegionRequest {
114    pub region_id: RegionId,
115    pub table_dir: String,
116    pub path_type: PathType,
117    pub region_options: RegionOptions,
118    pub max_parallelism: usize,
119}
120
121/// Open a compaction region from a compaction request.
122/// It's simple version of RegionOpener::open().
123pub async fn open_compaction_region(
124    req: &OpenCompactionRegionRequest,
125    mito_config: &MitoConfig,
126    object_store_manager: ObjectStoreManagerRef,
127    ttl_provider: Either<TimeToLive, SchemaMetadataManagerRef>,
128) -> Result<CompactionRegion> {
129    let object_store = {
130        let name = &req.region_options.storage;
131        if let Some(name) = name {
132            object_store_manager
133                .find(name)
134                .with_context(|| ObjectStoreNotFoundSnafu {
135                    object_store: name.clone(),
136                })?
137        } else {
138            object_store_manager.default_object_store()
139        }
140    };
141
142    let access_layer = {
143        let puffin_manager_factory = PuffinManagerFactory::new(
144            &mito_config.index.aux_path,
145            mito_config.index.staging_size.as_bytes(),
146            Some(mito_config.index.write_buffer_size.as_bytes() as _),
147            mito_config.index.staging_ttl,
148        )
149        .await?;
150        let intermediate_manager =
151            IntermediateManager::init_fs(mito_config.index.aux_path.clone()).await?;
152
153        Arc::new(AccessLayer::new(
154            &req.table_dir,
155            req.path_type,
156            object_store.clone(),
157            puffin_manager_factory,
158            intermediate_manager,
159        ))
160    };
161
162    let manifest_manager = {
163        let region_dir = region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type);
164        let region_manifest_options =
165            RegionManifestOptions::new(mito_config, &region_dir, object_store);
166
167        RegionManifestManager::open(region_manifest_options, &Default::default())
168            .await?
169            .with_context(|| EmptyRegionDirSnafu {
170                region_id: req.region_id,
171                region_dir: region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type),
172            })?
173    };
174
175    let manifest = manifest_manager.manifest();
176    let region_metadata = manifest.metadata.clone();
177    let manifest_ctx = Arc::new(ManifestContext::new(
178        manifest_manager,
179        RegionRoleState::Leader(RegionLeaderState::Writable),
180    ));
181
182    let file_purger = {
183        let purge_scheduler = Arc::new(LocalScheduler::new(mito_config.max_background_purges));
184        Arc::new(LocalFilePurger::new(
185            purge_scheduler.clone(),
186            access_layer.clone(),
187            None,
188        ))
189    };
190
191    let current_version = {
192        let mut ssts = SstVersion::new();
193        ssts.add_files(file_purger.clone(), manifest.files.values().cloned());
194        CompactionVersion {
195            metadata: region_metadata.clone(),
196            options: req.region_options.clone(),
197            ssts: Arc::new(ssts),
198            compaction_time_window: manifest.compaction_time_window,
199        }
200    };
201
202    let ttl = match ttl_provider {
203        // Use the specified ttl.
204        Either::Left(ttl) => ttl,
205        // Get the ttl from the schema metadata manager.
206        Either::Right(schema_metadata_manager) => {
207            let (_, ttl) = find_dynamic_options(
208                req.region_id.table_id(),
209                &req.region_options,
210                &schema_metadata_manager,
211            )
212            .await
213            .unwrap_or_else(|e| {
214                warn!(e; "Failed to get ttl for region: {}", region_metadata.region_id);
215                (
216                    crate::region::options::CompactionOptions::default(),
217                    TimeToLive::default(),
218                )
219            });
220            ttl
221        }
222    };
223
224    Ok(CompactionRegion {
225        region_id: req.region_id,
226        region_options: req.region_options.clone(),
227        engine_config: Arc::new(mito_config.clone()),
228        region_metadata: region_metadata.clone(),
229        cache_manager: Arc::new(CacheManager::default()),
230        access_layer,
231        manifest_ctx,
232        current_version,
233        file_purger: Some(file_purger),
234        ttl: Some(ttl),
235        max_parallelism: req.max_parallelism,
236    })
237}
238
239impl CompactionRegion {
240    /// Get the file purger of the compaction region.
241    pub fn file_purger(&self) -> Option<Arc<LocalFilePurger>> {
242        self.file_purger.clone()
243    }
244
245    /// Stop the file purger scheduler of the compaction region.
246    pub async fn stop_purger_scheduler(&self) -> Result<()> {
247        if let Some(file_purger) = &self.file_purger {
248            file_purger.stop_scheduler().await
249        } else {
250            Ok(())
251        }
252    }
253}
254
255/// `[MergeOutput]` represents the output of merging SST files.
256#[derive(Default, Clone, Debug, Serialize, Deserialize)]
257pub struct MergeOutput {
258    pub files_to_add: Vec<FileMeta>,
259    pub files_to_remove: Vec<FileMeta>,
260    pub compaction_time_window: Option<i64>,
261}
262
263impl MergeOutput {
264    pub fn is_empty(&self) -> bool {
265        self.files_to_add.is_empty() && self.files_to_remove.is_empty()
266    }
267
268    pub fn input_file_size(&self) -> u64 {
269        self.files_to_remove.iter().map(|f| f.file_size).sum()
270    }
271
272    pub fn output_file_size(&self) -> u64 {
273        self.files_to_add.iter().map(|f| f.file_size).sum()
274    }
275}
276
277/// Compactor is the trait that defines the compaction logic.
278#[async_trait::async_trait]
279pub trait Compactor: Send + Sync + 'static {
280    /// Merge SST files for a region.
281    async fn merge_ssts(
282        &self,
283        compaction_region: &CompactionRegion,
284        picker_output: PickerOutput,
285    ) -> Result<MergeOutput>;
286
287    /// Update the manifest after merging SST files.
288    async fn update_manifest(
289        &self,
290        compaction_region: &CompactionRegion,
291        merge_output: MergeOutput,
292    ) -> Result<RegionEdit>;
293
294    /// Execute compaction for a region.
295    async fn compact(
296        &self,
297        compaction_region: &CompactionRegion,
298        compact_request_options: compact_request::Options,
299    ) -> Result<()>;
300}
301
302/// DefaultCompactor is the default implementation of Compactor.
303pub struct DefaultCompactor;
304
305impl DefaultCompactor {
306    /// Merge a single compaction output into SST files.
307    async fn merge_single_output(
308        compaction_region: CompactionRegion,
309        output: CompactionOutput,
310        write_opts: WriteOptions,
311    ) -> Result<Vec<FileMeta>> {
312        let region_id = compaction_region.region_id;
313        let storage = compaction_region.region_options.storage.clone();
314        let index_options = compaction_region
315            .current_version
316            .options
317            .index_options
318            .clone();
319        let append_mode = compaction_region.current_version.options.append_mode;
320        let merge_mode = compaction_region.current_version.options.merge_mode();
321        let flat_format = compaction_region
322            .region_options
323            .sst_format
324            .map(|format| format == FormatType::Flat)
325            .unwrap_or(compaction_region.engine_config.default_flat_format);
326
327        let index_config = compaction_region.engine_config.index.clone();
328        let inverted_index_config = compaction_region.engine_config.inverted_index.clone();
329        let fulltext_index_config = compaction_region.engine_config.fulltext_index.clone();
330        let bloom_filter_index_config = compaction_region.engine_config.bloom_filter_index.clone();
331        #[cfg(feature = "vector_index")]
332        let vector_index_config = compaction_region.engine_config.vector_index.clone();
333
334        let input_file_names = output
335            .inputs
336            .iter()
337            .map(|f| f.file_id().to_string())
338            .join(",");
339        let max_sequence = output
340            .inputs
341            .iter()
342            .map(|f| f.meta_ref().sequence)
343            .max()
344            .flatten();
345        let builder = CompactionSstReaderBuilder {
346            metadata: compaction_region.region_metadata.clone(),
347            sst_layer: compaction_region.access_layer.clone(),
348            cache: compaction_region.cache_manager.clone(),
349            inputs: &output.inputs,
350            append_mode,
351            filter_deleted: output.filter_deleted,
352            time_range: output.output_time_range,
353            merge_mode,
354        };
355        let reader = builder.build_flat_sst_reader().await?;
356        let source = FlatSource::Stream(reader);
357        let mut metrics = Metrics::new(WriteType::Compaction);
358        let region_metadata = compaction_region.region_metadata.clone();
359        let sst_infos = compaction_region
360            .access_layer
361            .write_sst(
362                SstWriteRequest {
363                    op_type: OperationType::Compact,
364                    metadata: region_metadata.clone(),
365                    source,
366                    cache_manager: compaction_region.cache_manager.clone(),
367                    storage,
368                    max_sequence: max_sequence.map(NonZero::get),
369                    sst_write_format: if flat_format {
370                        FormatType::Flat
371                    } else {
372                        FormatType::PrimaryKey
373                    },
374                    index_options,
375                    index_config,
376                    inverted_index_config,
377                    fulltext_index_config,
378                    bloom_filter_index_config,
379                    #[cfg(feature = "vector_index")]
380                    vector_index_config,
381                },
382                &write_opts,
383                &mut metrics,
384            )
385            .await?;
386        // Convert partition expression once outside the map
387        let partition_expr = match &region_metadata.partition_expr {
388            None => None,
389            Some(json_str) if json_str.is_empty() => None,
390            Some(json_str) => PartitionExpr::from_json_str(json_str).with_context(|_| {
391                InvalidPartitionExprSnafu {
392                    expr: json_str.clone(),
393                }
394            })?,
395        };
396
397        let output_files = sst_infos
398            .into_iter()
399            .map(|sst_info| FileMeta {
400                region_id,
401                file_id: sst_info.file_id,
402                time_range: sst_info.time_range,
403                level: output.output_level,
404                file_size: sst_info.file_size,
405                max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
406                available_indexes: sst_info.index_metadata.build_available_indexes(),
407                indexes: sst_info.index_metadata.build_indexes(),
408                index_file_size: sst_info.index_metadata.file_size,
409                index_version: 0,
410                num_rows: sst_info.num_rows as u64,
411                num_row_groups: sst_info.num_row_groups,
412                sequence: max_sequence,
413                partition_expr: partition_expr.clone(),
414                num_series: sst_info.num_series,
415            })
416            .collect::<Vec<_>>();
417        let output_file_names = output_files.iter().map(|f| f.file_id.to_string()).join(",");
418        info!(
419            "Region {} compaction inputs: [{}], outputs: [{}], flat_format: {}, metrics: {:?}",
420            region_id, input_file_names, output_file_names, flat_format, metrics
421        );
422        metrics.observe();
423        Ok(output_files)
424    }
425}
426
427#[async_trait::async_trait]
428impl Compactor for DefaultCompactor {
429    async fn merge_ssts(
430        &self,
431        compaction_region: &CompactionRegion,
432        mut picker_output: PickerOutput,
433    ) -> Result<MergeOutput> {
434        let mut futs = Vec::with_capacity(picker_output.outputs.len());
435        let mut compacted_inputs =
436            Vec::with_capacity(picker_output.outputs.iter().map(|o| o.inputs.len()).sum());
437        let internal_parallelism = compaction_region.max_parallelism.max(1);
438        let compaction_time_window = picker_output.time_window_size;
439
440        for output in picker_output.outputs.drain(..) {
441            let inputs_to_remove: Vec<_> =
442                output.inputs.iter().map(|f| f.meta_ref().clone()).collect();
443            compacted_inputs.extend(inputs_to_remove.iter().cloned());
444            let write_opts = WriteOptions {
445                write_buffer_size: compaction_region.engine_config.sst_write_buffer_size,
446                max_file_size: picker_output.max_file_size,
447                ..Default::default()
448            };
449            futs.push(Self::merge_single_output(
450                compaction_region.clone(),
451                output,
452                write_opts,
453            ));
454        }
455        let mut output_files = Vec::with_capacity(futs.len());
456        while !futs.is_empty() {
457            let mut task_chunk = Vec::with_capacity(internal_parallelism);
458            for _ in 0..internal_parallelism {
459                if let Some(task) = futs.pop() {
460                    task_chunk.push(common_runtime::spawn_compact(task));
461                }
462            }
463            let metas = futures::future::try_join_all(task_chunk)
464                .await
465                .context(JoinSnafu)?
466                .into_iter()
467                .collect::<Result<Vec<Vec<_>>>>()?;
468            output_files.extend(metas.into_iter().flatten());
469        }
470
471        // In case of remote compaction, we still allow the region edit after merge to
472        // clean expired ssts.
473        let mut inputs: Vec<_> = compacted_inputs.into_iter().collect();
474        inputs.extend(
475            picker_output
476                .expired_ssts
477                .iter()
478                .map(|f| f.meta_ref().clone()),
479        );
480
481        Ok(MergeOutput {
482            files_to_add: output_files,
483            files_to_remove: inputs,
484            compaction_time_window: Some(compaction_time_window),
485        })
486    }
487
488    async fn update_manifest(
489        &self,
490        compaction_region: &CompactionRegion,
491        merge_output: MergeOutput,
492    ) -> Result<RegionEdit> {
493        // Write region edit to manifest.
494        let edit = RegionEdit {
495            files_to_add: merge_output.files_to_add,
496            files_to_remove: merge_output.files_to_remove,
497            // Use current timestamp as the edit timestamp.
498            timestamp_ms: Some(chrono::Utc::now().timestamp_millis()),
499            compaction_time_window: merge_output
500                .compaction_time_window
501                .map(|seconds| Duration::from_secs(seconds as u64)),
502            flushed_entry_id: None,
503            flushed_sequence: None,
504            committed_sequence: None,
505        };
506
507        let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone()));
508        // TODO: We might leak files if we fail to update manifest. We can add a cleanup task to remove them later.
509        compaction_region
510            .manifest_ctx
511            .update_manifest(RegionLeaderState::Writable, action_list, false)
512            .await?;
513
514        Ok(edit)
515    }
516
517    // The default implementation of compact combines the merge_ssts and update_manifest functions.
518    // Note: It's local compaction and only used for testing purpose.
519    async fn compact(
520        &self,
521        compaction_region: &CompactionRegion,
522        compact_request_options: compact_request::Options,
523    ) -> Result<()> {
524        let picker_output = {
525            let picker_output = new_picker(
526                &compact_request_options,
527                &compaction_region.region_options.compaction,
528                compaction_region.region_options.append_mode,
529                None,
530            )
531            .pick(compaction_region);
532
533            if let Some(picker_output) = picker_output {
534                picker_output
535            } else {
536                info!(
537                    "No files to compact for region_id: {}",
538                    compaction_region.region_id
539                );
540                return Ok(());
541            }
542        };
543
544        let merge_output = self.merge_ssts(compaction_region, picker_output).await?;
545        if merge_output.is_empty() {
546            info!(
547                "No files to compact for region_id: {}",
548                compaction_region.region_id
549            );
550            return Ok(());
551        }
552
553        metrics::COMPACTION_INPUT_BYTES.inc_by(merge_output.input_file_size() as f64);
554        metrics::COMPACTION_OUTPUT_BYTES.inc_by(merge_output.output_file_size() as f64);
555        self.update_manifest(compaction_region, merge_output)
556            .await?;
557
558        Ok(())
559    }
560}