mito2/compaction/
reader.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};
16use std::sync::Arc;
17
18use common_time::Timestamp;
19use common_time::range::TimestampRange;
20use common_time::timestamp::TimeUnit;
21use datafusion_common::ScalarValue;
22use datafusion_expr::Expr;
23use datatypes::extension::json::is_json2_extension_type;
24use datatypes::types::json_type::JsonNativeType;
25use parquet::arrow::parquet_to_arrow_schema;
26use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
27use snafu::{OptionExt, ResultExt};
28use store_api::metadata::RegionMetadataRef;
29
30use crate::access_layer::AccessLayerRef;
31use crate::cache::{CacheManagerRef, CacheStrategy};
32use crate::error::{
33    DataTypeMismatchSnafu, ParquetToArrowSchemaSnafu, Result, TimeRangePredicateOverflowSnafu,
34};
35use crate::read::FlatSource;
36use crate::read::flat_projection::FlatProjectionMapper;
37use crate::read::read_columns::ReadColumns;
38use crate::read::scan_region::{PredicateGroup, ScanInput};
39use crate::read::seq_scan::SeqScan;
40use crate::region::options::MergeMode;
41use crate::sst::file::FileHandle;
42use crate::sst::parquet::reader::MetadataCacheMetrics;
43
44/// Builders to create [BoxedRecordBatchStream] for compaction.
45pub(crate) struct CompactionSstReaderBuilder<'a> {
46    pub(crate) metadata: RegionMetadataRef,
47    pub(crate) sst_layer: AccessLayerRef,
48    pub(crate) cache: CacheManagerRef,
49    pub(crate) inputs: &'a [FileHandle],
50    pub(crate) append_mode: bool,
51    pub(crate) filter_deleted: bool,
52    pub(crate) time_range: Option<TimestampRange>,
53    pub(crate) merge_mode: MergeMode,
54}
55
56impl CompactionSstReaderBuilder<'_> {
57    /// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files,
58    /// for compaction. The schema of the [FlatSource] is unified.
59    pub(crate) async fn build_flat_sst_reader(self) -> Result<FlatSource> {
60        let scan_input = self.build_scan_input().await?;
61
62        let schema = scan_input.mapper.output_schema();
63        let schema = schema.arrow_schema();
64
65        let stream = SeqScan::new(scan_input)
66            .build_flat_reader_for_compaction()
67            .await?;
68        Ok(FlatSource::new_stream(schema.clone(), stream))
69    }
70
71    async fn build_scan_input(self) -> Result<ScanInput> {
72        let schema = self.metadata.schema.arrow_schema();
73        let parquet_metadata = self.collect_parquet_metadata().await?;
74        let batch_size = crate::batch_size::estimate_batch_size(
75            parquet_metadata
76                .iter()
77                .flat_map(|metadata| metadata.row_groups())
78                .map(|row_group| {
79                    let uncompressed_bytes = row_group
80                        .columns()
81                        .iter()
82                        .map(|column| column.uncompressed_size() as u64)
83                        .sum();
84                    (row_group.num_rows() as u64, uncompressed_bytes)
85                }),
86        );
87        let json_type_hint = if schema.fields().iter().any(is_json2_extension_type) {
88            let mut json_type_hint = schema
89                .fields()
90                .iter()
91                .filter(|&field| is_json2_extension_type(field))
92                .map(|field| (field.name().clone(), JsonNativeType::Null))
93                .collect::<HashMap<_, _>>();
94
95            for metadata in &parquet_metadata {
96                let file_metadata = metadata.file_metadata();
97                let schema = parquet_to_arrow_schema(
98                    file_metadata.schema_descr(),
99                    file_metadata.key_value_metadata(),
100                )
101                .context(ParquetToArrowSchemaSnafu {
102                    file: "compaction input",
103                })?;
104                for field in schema.fields() {
105                    let Some(merged) = json_type_hint.get_mut(field.name()) else {
106                        continue;
107                    };
108
109                    let json_type = JsonNativeType::try_from(field.data_type())
110                        .context(DataTypeMismatchSnafu)?;
111                    merged.merge(&json_type);
112                }
113            }
114
115            Some(json_type_hint)
116        } else {
117            None
118        };
119
120        let projection = (0..self.metadata.column_metadatas.len()).collect();
121        let read_column_ids = self
122            .metadata
123            .column_metadatas
124            .iter()
125            .map(|x| x.column_id)
126            .collect::<Vec<_>>();
127        let json_target_types = json_type_hint
128            .as_ref()
129            .map(|hint| {
130                hint.iter()
131                    .filter_map(|(col_name, json_type)| {
132                        self.metadata
133                            .column_by_name(col_name)
134                            .map(|col| (col.column_id, json_type.clone()))
135                    })
136                    .collect::<BTreeMap<_, _>>()
137            })
138            .unwrap_or_default();
139        let read_columns =
140            ReadColumns::new(read_column_ids).with_json_target_types(json_target_types);
141        let mapper =
142            FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?;
143
144        let mut scan_input = ScanInput::new(self.sst_layer, mapper)
145            .with_files(self.inputs.to_vec())
146            .with_compaction(true)
147            .with_batch_size(batch_size)
148            .with_append_mode(self.append_mode)
149            // We use special cache strategy for compaction.
150            .with_cache(CacheStrategy::Compaction(self.cache))
151            .with_filter_deleted(self.filter_deleted)
152            // We ignore file not found error during compaction.
153            .with_ignore_file_not_found(true)
154            .with_merge_mode(self.merge_mode);
155
156        // This serves as a workaround of https://github.com/GreptimeTeam/greptimedb/issues/3944
157        // by converting time ranges into predicate.
158        if let Some(time_range) = self.time_range {
159            scan_input =
160                scan_input.with_predicate(time_range_to_predicate(time_range, &self.metadata)?);
161        }
162
163        Ok(scan_input)
164    }
165
166    async fn collect_parquet_metadata(&self) -> Result<Vec<Arc<ParquetMetaData>>> {
167        let mut metadata = Vec::with_capacity(self.inputs.len());
168
169        for file_handle in self.inputs {
170            let file_path =
171                file_handle.file_path(self.sst_layer.table_dir(), self.sst_layer.path_type());
172            let file_size = file_handle.meta_ref().file_size;
173            let parquet_metadata = match self
174                .sst_layer
175                .read_sst(file_handle.clone())
176                .cache(CacheStrategy::Compaction(self.cache.clone()))
177                .read_parquet_metadata(
178                    &file_path,
179                    file_size,
180                    &mut MetadataCacheMetrics::default(),
181                    PageIndexPolicy::default(),
182                )
183                .await
184                .map(|x| x.0.parquet_metadata())
185            {
186                Ok(x) => x,
187                Err(e) if e.is_object_not_found() => continue,
188                Err(e) => return Err(e),
189            };
190            metadata.push(parquet_metadata);
191        }
192        Ok(metadata)
193    }
194}
195
196/// Converts time range to predicates so that rows outside the range will be filtered.
197fn time_range_to_predicate(
198    range: TimestampRange,
199    metadata: &RegionMetadataRef,
200) -> Result<PredicateGroup> {
201    let ts_col = metadata.time_index_column();
202
203    // safety: time index column's type must be a valid timestamp type.
204    let ts_col_unit = ts_col
205        .column_schema
206        .data_type
207        .as_timestamp()
208        .unwrap()
209        .unit();
210
211    let exprs = match (range.start(), range.end()) {
212        (Some(start), Some(end)) => {
213            vec![
214                datafusion_expr::col(ts_col.column_schema.name.clone())
215                    .gt_eq(ts_to_lit(*start, ts_col_unit)?),
216                datafusion_expr::col(ts_col.column_schema.name.clone())
217                    .lt(ts_to_lit(*end, ts_col_unit)?),
218            ]
219        }
220        (Some(start), None) => {
221            vec![
222                datafusion_expr::col(ts_col.column_schema.name.clone())
223                    .gt_eq(ts_to_lit(*start, ts_col_unit)?),
224            ]
225        }
226
227        (None, Some(end)) => {
228            vec![
229                datafusion_expr::col(ts_col.column_schema.name.clone())
230                    .lt(ts_to_lit(*end, ts_col_unit)?),
231            ]
232        }
233        (None, None) => {
234            return Ok(PredicateGroup::default());
235        }
236    };
237
238    let predicate = PredicateGroup::new(metadata, &exprs)?;
239    Ok(predicate)
240}
241
242fn ts_to_lit(ts: Timestamp, ts_col_unit: TimeUnit) -> Result<Expr> {
243    let ts = ts
244        .convert_to(ts_col_unit)
245        .context(TimeRangePredicateOverflowSnafu {
246            timestamp: ts,
247            unit: ts_col_unit,
248        })?;
249    let val = ts.value();
250    let scalar_value = match ts_col_unit {
251        TimeUnit::Second => ScalarValue::TimestampSecond(Some(val), None),
252        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(val), None),
253        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(val), None),
254        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(val), None),
255    };
256    Ok(datafusion_expr::lit(scalar_value))
257}
mito2/compaction/reader.rs

mito2/compaction/
reader.rs
