mito2/memtable/bulk/
row_group_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::sync::Arc;
16
17use bytes::Bytes;
18use datatypes::extension::json::is_structured_json_field;
19use parquet::arrow::ProjectionMask;
20use parquet::arrow::arrow_reader::{
21    ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader,
22    ParquetRecordBatchReaderBuilder, RowSelection,
23};
24use parquet::file::metadata::ParquetMetaData;
25use snafu::ResultExt;
26
27use crate::error;
28use crate::error::ReadDataPartSnafu;
29use crate::memtable::bulk::chunk_reader::MemtableChunkReader;
30use crate::memtable::bulk::context::BulkIterContextRef;
31use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
32
33pub(crate) struct MemtableRowGroupReaderBuilder {
34    projection: ProjectionMask,
35    arrow_metadata: ArrowReaderMetadata,
36    data: Bytes,
37}
38
39impl MemtableRowGroupReaderBuilder {
40    pub(crate) fn try_new(
41        context: &BulkIterContextRef,
42        projection: ProjectionMask,
43        parquet_metadata: Arc<ParquetMetaData>,
44        data: Bytes,
45    ) -> error::Result<Self> {
46        // Create ArrowReaderMetadata for building the reader.
47        let mut arrow_reader_options = ArrowReaderOptions::new();
48
49        // JSON2 columns in region metadata are not concretized with nested
50        // fields here, so let parquet use its embedded Arrow schema instead.
51        //
52        // TODO: Pass the encoded part's concrete schema into this builder. It
53        // avoids parsing the Arrow schema from parquet metadata while keeping
54        // JSON2 nested fields exact.
55        if !context
56            .read_format()
57            .arrow_schema()
58            .fields()
59            .iter()
60            .any(is_structured_json_field)
61        {
62            arrow_reader_options =
63                arrow_reader_options.with_schema(context.read_format().arrow_schema().clone());
64        }
65
66        let arrow_metadata =
67            ArrowReaderMetadata::try_new(parquet_metadata.clone(), arrow_reader_options)
68                .context(ReadDataPartSnafu)?;
69        Ok(Self {
70            projection,
71            arrow_metadata,
72            data,
73        })
74    }
75
76    /// Builds a reader to read the row group at `row_group_idx` from memory.
77    pub(crate) fn build_row_group_reader(
78        &self,
79        row_group_idx: usize,
80        row_selection: Option<RowSelection>,
81    ) -> error::Result<ParquetRecordBatchReader> {
82        let chunk_reader = MemtableChunkReader::new(self.data.clone());
83
84        let mut builder = ParquetRecordBatchReaderBuilder::new_with_metadata(
85            chunk_reader,
86            self.arrow_metadata.clone(),
87        )
88        .with_row_groups(vec![row_group_idx])
89        .with_projection(self.projection.clone())
90        .with_batch_size(DEFAULT_READ_BATCH_SIZE);
91
92        if let Some(selection) = row_selection {
93            builder = builder.with_row_selection(selection);
94        }
95
96        builder.build().context(ReadDataPartSnafu)
97    }
98}
mito2/memtable/bulk/row_group_reader.rs

mito2/memtable/bulk/
row_group_reader.rs
