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

mito2/compaction/
reader.rs
