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;
16use std::sync::Arc;
17
18use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY;
19use common_time::Timestamp;
20use common_time::range::TimestampRange;
21use common_time::timestamp::TimeUnit;
22use datafusion_common::ScalarValue;
23use datafusion_expr::Expr;
24use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
25use snafu::OptionExt;
26use store_api::metadata::RegionMetadataRef;
27
28use crate::access_layer::AccessLayerRef;
29use crate::cache::{CacheManagerRef, CacheStrategy};
30use crate::compaction::json2::{
31    Json2RewritePlans, collect_json2_rewrite_plans_from_parquet, rewrite_json2_schema,
32};
33use crate::error::{InvalidRecordBatchSnafu, Result, TimeRangePredicateOverflowSnafu};
34use crate::read::FlatSource;
35use crate::read::flat_projection::FlatProjectionMapper;
36use crate::read::read_columns::ReadColumns;
37use crate::read::scan_region::{PredicateGroup, ScanInput};
38use crate::read::seq_scan::SeqScan;
39use crate::read::stream::spawn_compact_stream;
40use crate::region::options::MergeMode;
41use crate::sst::file::FileHandle;
42use crate::sst::parquet::reader::MetadataCacheMetrics;
43use crate::sst::parquet::{Json2RewriteTargets, Json2TargetLayout};
44
45/// Builders to create [BoxedRecordBatchStream] for compaction.
46pub(crate) struct CompactionSstReaderBuilder<'a> {
47    pub(crate) metadata: RegionMetadataRef,
48    pub(crate) sst_layer: AccessLayerRef,
49    pub(crate) cache: CacheManagerRef,
50    pub(crate) inputs: &'a [FileHandle],
51    pub(crate) append_mode: bool,
52    pub(crate) filter_deleted: bool,
53    pub(crate) time_range: Option<TimestampRange>,
54    pub(crate) merge_mode: MergeMode,
55    /// Max number of inputs decoded concurrently ahead of the merge.
56    pub(crate) read_parallelism: usize,
57}
58
59impl CompactionSstReaderBuilder<'_> {
60    /// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files,
61    /// for compaction. The schema of the [FlatSource] is unified.
62    pub(crate) async fn build_flat_sst_reader(self) -> Result<FlatSource> {
63        let parquet_metadata = self.collect_parquet_metadata().await?;
64        let plans = collect_json2_rewrite_plans_from_parquet(&self.metadata, &parquet_metadata)?;
65        let scan_input = self.build_scan_input(&parquet_metadata, &plans)?;
66
67        let schema = scan_input.mapper.output_schema();
68        let schema = rewrite_json2_schema(schema.arrow_schema(), &plans);
69
70        let stream = SeqScan::new(scan_input)
71            .build_flat_reader_for_compaction()
72            .await?;
73        Ok(FlatSource::new_stream(schema, stream))
74    }
75
76    fn build_scan_input(
77        self,
78        parquet_metadata: &[Arc<ParquetMetaData>],
79        plans: &Json2RewritePlans,
80    ) -> Result<ScanInput> {
81        let batch_size = crate::batch_size::estimate_batch_size(
82            parquet_metadata
83                .iter()
84                .flat_map(|metadata| metadata.row_groups())
85                .map(|row_group| {
86                    let uncompressed_bytes = row_group
87                        .columns()
88                        .iter()
89                        .map(|column| column.uncompressed_size() as u64)
90                        .sum();
91                    (row_group.num_rows() as u64, uncompressed_bytes)
92                }),
93        );
94
95        let projection = (0..self.metadata.column_metadatas.len()).collect();
96        let read_column_ids = self
97            .metadata
98            .column_metadatas
99            .iter()
100            .map(|x| x.column_id)
101            .collect::<Vec<_>>();
102
103        let mut json2_target_layouts = BTreeMap::new();
104        for (name, plan) in plans {
105            let Some(column) = self.metadata.column_by_name(name) else {
106                continue;
107            };
108            let extension_metadata = column
109                .column_schema
110                .metadata()
111                .get(EXTENSION_TYPE_METADATA_KEY)
112                .cloned()
113                .with_context(|| InvalidRecordBatchSnafu {
114                    reason: format!("JSON2 target column '{name}' has no extension metadata"),
115                })?;
116            json2_target_layouts.insert(
117                column.column_id,
118                Json2TargetLayout {
119                    extension_metadata,
120                    target_layout: plan.target_layout.clone(),
121                },
122            );
123        }
124        let read_columns = ReadColumns::new(read_column_ids);
125        let targets: Json2RewriteTargets = Arc::new(json2_target_layouts);
126        let mapper = FlatProjectionMapper::new_with_json2_rewrite_targets(
127            &self.metadata,
128            projection,
129            read_columns,
130            &targets,
131        )?;
132
133        let mut scan_input = ScanInput::builder(self.sst_layer, mapper)
134            .with_json2_rewrite_targets(targets)
135            .with_files(self.inputs.to_vec())
136            .with_compaction(true)
137            .with_compaction_read_parallelism(self.read_parallelism)
138            .with_batch_size(batch_size)
139            .with_append_mode(self.append_mode)
140            // We use special cache strategy for compaction.
141            .with_cache(CacheStrategy::Compaction(self.cache))
142            .with_filter_deleted(self.filter_deleted)
143            // We ignore file not found error during compaction.
144            .with_ignore_file_not_found(true)
145            .with_merge_mode(self.merge_mode);
146
147        // This serves as a workaround of https://github.com/GreptimeTeam/greptimedb/issues/3944
148        // by converting time ranges into predicate.
149        if let Some(time_range) = self.time_range {
150            scan_input =
151                scan_input.with_predicate(time_range_to_predicate(time_range, &self.metadata)?);
152        }
153
154        Ok(scan_input.build())
155    }
156
157    async fn collect_parquet_metadata(&self) -> Result<Vec<Arc<ParquetMetaData>>> {
158        let mut metadata = Vec::with_capacity(self.inputs.len());
159
160        for file_handle in self.inputs {
161            let file_path =
162                file_handle.file_path(self.sst_layer.table_dir(), self.sst_layer.path_type());
163            let file_size = file_handle.meta_ref().file_size;
164            let parquet_metadata = match self
165                .sst_layer
166                .read_sst(file_handle.clone())
167                .cache(CacheStrategy::Compaction(self.cache.clone()))
168                .read_parquet_metadata(
169                    &file_path,
170                    file_size,
171                    &mut MetadataCacheMetrics::default(),
172                    PageIndexPolicy::default(),
173                )
174                .await
175                .map(|x| x.0.parquet_metadata())
176            {
177                Ok(x) => x,
178                Err(e) if e.is_object_not_found() => continue,
179                Err(e) => return Err(e),
180            };
181            metadata.push(parquet_metadata);
182        }
183        Ok(metadata)
184    }
185}
186
187/// Number of merged batches buffered between the merge task and the SST writer.
188const MERGED_BATCH_CHANNEL_SIZE: usize = 2;
189
190/// Reads and merges `source` on a separate task, so the SST writer can encode
191/// merged batches while the next ones are being read and merged.
192pub(crate) fn spawn_merge(source: FlatSource) -> FlatSource {
193    let schema = source.schema().clone();
194    let stream = spawn_compact_stream(source.into_stream(), MERGED_BATCH_CHANNEL_SIZE, None);
195    FlatSource::new_stream(schema, stream)
196}
197
198/// Converts time range to predicates so that rows outside the range will be filtered.
199fn time_range_to_predicate(
200    range: TimestampRange,
201    metadata: &RegionMetadataRef,
202) -> Result<PredicateGroup> {
203    let ts_col = metadata.time_index_column();
204
205    // safety: time index column's type must be a valid timestamp type.
206    let ts_col_unit = ts_col
207        .column_schema
208        .data_type
209        .as_timestamp()
210        .unwrap()
211        .unit();
212
213    let exprs = match (range.start(), range.end()) {
214        (Some(start), Some(end)) => {
215            vec![
216                datafusion_expr::col(ts_col.column_schema.name.clone())
217                    .gt_eq(ts_to_lit(*start, ts_col_unit)?),
218                datafusion_expr::col(ts_col.column_schema.name.clone())
219                    .lt(ts_to_lit(*end, ts_col_unit)?),
220            ]
221        }
222        (Some(start), None) => {
223            vec![
224                datafusion_expr::col(ts_col.column_schema.name.clone())
225                    .gt_eq(ts_to_lit(*start, ts_col_unit)?),
226            ]
227        }
228
229        (None, Some(end)) => {
230            vec![
231                datafusion_expr::col(ts_col.column_schema.name.clone())
232                    .lt(ts_to_lit(*end, ts_col_unit)?),
233            ]
234        }
235        (None, None) => {
236            return Ok(PredicateGroup::default());
237        }
238    };
239
240    let predicate = PredicateGroup::new(metadata, &exprs)?;
241    Ok(predicate)
242}
243
244fn ts_to_lit(ts: Timestamp, ts_col_unit: TimeUnit) -> Result<Expr> {
245    let ts = ts
246        .convert_to(ts_col_unit)
247        .context(TimeRangePredicateOverflowSnafu {
248            timestamp: ts,
249            unit: ts_col_unit,
250        })?;
251    let val = ts.value();
252    let scalar_value = match ts_col_unit {
253        TimeUnit::Second => ScalarValue::TimestampSecond(Some(val), None),
254        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(val), None),
255        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(val), None),
256        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(val), None),
257    };
258    Ok(datafusion_expr::lit(scalar_value))
259}
260
261#[cfg(test)]
262mod tests {
263    use api::v1::Rows;
264    use store_api::region_engine::RegionEngine;
265    use store_api::region_request::{RegionFlushRequest, RegionRequest};
266    use store_api::storage::RegionId;
267
268    use super::*;
269    use crate::test_util::{
270        CreateRequestBuilder, TestEnv, build_rows_for_key, put_rows, rows_schema,
271    };
272
273    #[tokio::test]
274    async fn test_compaction_reader_keeps_short_series_batches_whole() {
275        let mut env = TestEnv::new().await;
276        let engine = env.create_engine(Default::default()).await;
277        let region_id = RegionId::new(1, 1);
278        let request = CreateRequestBuilder::new().build();
279        let schema = rows_schema(&request);
280        engine
281            .handle_request(region_id, RegionRequest::Create(request))
282            .await
283            .unwrap();
284        // Each file has 100 short series, which queries split for merging.
285        for file in 0..2 {
286            let rows = (0..100)
287                .flat_map(|series| {
288                    build_rows_for_key(&format!("key-{series:03}"), file * 10, file * 10 + 10, 0)
289                })
290                .collect();
291            put_rows(
292                &engine,
293                region_id,
294                Rows {
295                    schema: schema.clone(),
296                    rows,
297                },
298            )
299            .await;
300            engine
301                .handle_request(
302                    region_id,
303                    RegionRequest::Flush(RegionFlushRequest::default()),
304                )
305                .await
306                .unwrap();
307        }
308        let region = engine.get_region(region_id).unwrap();
309        let version = region.version();
310        let inputs: Vec<_> = version.ssts.levels()[0].files().cloned().collect();
311        assert_eq!(2, inputs.len());
312
313        let mut source = CompactionSstReaderBuilder {
314            metadata: version.metadata.clone(),
315            sst_layer: region.access_layer.clone(),
316            cache: engine.cache_manager(),
317            inputs: &inputs,
318            append_mode: false,
319            filter_deleted: true,
320            time_range: None,
321            merge_mode: MergeMode::LastRow,
322            read_parallelism: 1,
323        }
324        .build_flat_sst_reader()
325        .await
326        .unwrap();
327        let mut num_batches = 0;
328        let mut num_rows = 0;
329        while let Some(batch) = source.next_batch().await.unwrap() {
330            num_batches += 1;
331            num_rows += batch.num_rows();
332        }
333
334        // Splitting by series would yield a batch per series and file, 200 in total.
335        assert_eq!(2000, num_rows);
336        assert!(num_batches < 20, "got {num_batches} batches");
337    }
338}
mito2/compaction/reader.rs

mito2/compaction/
reader.rs
