risingwave_connector/parser/
parquet_parser.rs

1// Copyright 2024 RisingWave Labs
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 futures_async_stream::try_stream;
18use prometheus::core::GenericCounter;
19use risingwave_common::array::arrow::arrow_array_iceberg::{ArrayRef, RecordBatch};
20use risingwave_common::array::arrow::arrow_schema_iceberg::FieldRef;
21use risingwave_common::array::arrow::{IcebergArrowConvert, is_parquet_field_match_source_schema};
22use risingwave_common::array::{ArrayBuilderImpl, DataChunk, StreamChunk};
23use risingwave_common::metrics::LabelGuardedMetric;
24use risingwave_common::types::{Datum, ScalarImpl};
25use thiserror_ext::AsReport;
26
27use crate::parser::ConnectorResult;
28use crate::source::SourceColumnDesc;
29/// `ParquetParser` is responsible for converting the incoming `record_batch_stream`
30/// into a `streamChunk`.
31#[derive(Debug)]
32pub struct ParquetParser {
33    rw_columns: Vec<SourceColumnDesc>,
34    file_name: String,
35    offset: usize,
36    case_insensitive: bool,
37}
38
39impl ParquetParser {
40    pub fn new(
41        rw_columns: Vec<SourceColumnDesc>,
42        file_name: String,
43        offset: usize,
44        case_insensitive: bool,
45    ) -> ConnectorResult<Self> {
46        Ok(Self {
47            rw_columns,
48            file_name,
49            offset,
50            case_insensitive,
51        })
52    }
53
54    #[try_stream(boxed, ok = StreamChunk, error = crate::error::ConnectorError)]
55    pub async fn into_stream(
56        mut self,
57        record_batch_stream: parquet::arrow::async_reader::ParquetRecordBatchStream<
58            tokio_util::compat::Compat<opendal::FuturesAsyncReader>,
59        >,
60        file_source_input_row_count_metrics: Option<
61            LabelGuardedMetric<GenericCounter<prometheus::core::AtomicU64>>,
62        >,
63        parquet_source_skip_row_count_metrics: Option<
64            LabelGuardedMetric<GenericCounter<prometheus::core::AtomicU64>>,
65        >,
66    ) {
67        #[for_await]
68        for record_batch in record_batch_stream {
69            let record_batch: RecordBatch = record_batch?;
70            // Convert each record batch into a stream chunk according to user defined schema.
71            let chunk: StreamChunk = self.convert_record_batch_to_stream_chunk(
72                record_batch,
73                file_source_input_row_count_metrics.clone(),
74                parquet_source_skip_row_count_metrics.clone(),
75            )?;
76
77            yield chunk;
78        }
79    }
80
81    fn inc_offset(&mut self) {
82        self.offset += 1;
83    }
84
85    /// The function `convert_record_batch_to_stream_chunk` is designed to transform the given `RecordBatch` into a `StreamChunk`.
86    ///
87    /// For each column in the source column:
88    /// - If the column's schema matches a column in the `RecordBatch` (both the data type and column name are the same),
89    ///   the corresponding records are converted into a column of the `StreamChunk`.
90    /// - If the column's schema does not match, null values are inserted.
91    /// - Hidden columns are handled separately by filling in the appropriate fields to ensure the data chunk maintains the correct format.
92    /// - If a column in the Parquet file does not exist in the source schema, it is skipped.
93    ///
94    /// # Arguments
95    ///
96    /// * `record_batch` - The `RecordBatch` to be converted into a `StreamChunk`.
97    ///
98    /// # Returns
99    ///
100    /// A `StreamChunk` containing the converted data from the `RecordBatch`.
101    ///
102    /// The hidden columns that must be included here are `_rw_file` and `_rw_offset`.
103    /// Depending on whether the user specifies a primary key (pk), there may be an additional hidden column `row_id`.
104    /// Therefore, the maximum number of hidden columns is three.
105    fn convert_record_batch_to_stream_chunk(
106        &mut self,
107        record_batch: RecordBatch,
108        file_source_input_row_count_metrics: Option<
109            LabelGuardedMetric<GenericCounter<prometheus::core::AtomicU64>>,
110        >,
111        parquet_source_skip_row_count_metrics: Option<
112            LabelGuardedMetric<GenericCounter<prometheus::core::AtomicU64>>,
113        >,
114    ) -> Result<StreamChunk, crate::error::ConnectorError> {
115        const MAX_HIDDEN_COLUMN_NUMS: usize = 3;
116        let column_size = self.rw_columns.len();
117        let mut chunk_columns = Vec::with_capacity(self.rw_columns.len() + MAX_HIDDEN_COLUMN_NUMS);
118
119        for source_column in self.rw_columns.clone() {
120            match source_column.column_type {
121                crate::source::SourceColumnType::Normal => {
122                    let rw_data_type: &risingwave_common::types::DataType =
123                        &source_column.data_type;
124                    let rw_column_name = &source_column.name;
125                    if let Some((parquet_field, parquet_column)) =
126                        self.find_parquet_column(&record_batch, rw_column_name)
127                        && is_parquet_field_match_source_schema(parquet_field, rw_data_type)
128                    {
129                        // The match guard above verified this column converts to the
130                        // declared type; decode by the declared-side field, which for
131                        // variant carries the extension name that `from_array` routes on.
132                        let arrow_field = IcebergArrowConvert
133                            .to_arrow_field(rw_column_name, rw_data_type)
134                            .map_err(|e| {
135                                crate::parser::AccessError::ParquetParser {
136                                    message: format!(
137                                        "to_arrow_field failed, column='{}', rw_type='{}', offset={}, error={}",
138                                        rw_column_name, rw_data_type, self.offset, e.as_report()
139                                    )
140                                }
141                            })?;
142                        let array_impl = IcebergArrowConvert
143                            .array_from_arrow_array(&arrow_field, parquet_column)
144                            .map_err(|e| {
145                                crate::parser::AccessError::ParquetParser {
146                                    message: format!(
147                                        "array_from_arrow_array failed, column='{}', rw_type='{}', arrow_field='{}', parquet_type='{}', offset={}, error={}",
148                                        rw_column_name,
149                                        rw_data_type,
150                                        arrow_field.data_type(),
151                                        parquet_column.data_type(),
152                                        self.offset,
153                                        e.as_report()
154                                    )
155                                }
156                            })?;
157                        // The schema match and the decode are maintained in parallel; a
158                        // diverging output type is a code bug — surface it instead of letting
159                        // a mistyped column corrupt the chunk downstream.
160                        if array_impl.data_type() != *rw_data_type {
161                            return Err(crate::parser::AccessError::ParquetParser {
162                                message: format!(
163                                    "converted array type diverges from the declared column type, column='{}', rw_type='{}', converted='{}', parquet_type='{}', offset={}",
164                                    rw_column_name,
165                                    rw_data_type,
166                                    array_impl.data_type(),
167                                    parquet_column.data_type(),
168                                    self.offset,
169                                ),
170                            }
171                            .into());
172                        }
173                        chunk_columns.push(Arc::new(array_impl));
174                    } else {
175                        // Handle additional columns, for file source, the additional columns are offset and file name;
176                        // for columns defined in the user schema but not present in the parquet file, fill with null.
177                        let column = if let Some(additional_column_type) =
178                            &source_column.additional_column.column_type
179                        {
180                            match additional_column_type {
181                                risingwave_pb::plan_common::additional_column::ColumnType::Offset(_) => {
182                                    let mut array_builder = ArrayBuilderImpl::with_type(column_size, source_column.data_type.clone());
183                                    for _ in 0..record_batch.num_rows() {
184                                        let datum: Datum = Some(ScalarImpl::Utf8((self.offset).to_string().into()));
185                                        self.inc_offset();
186                                        array_builder.append(datum);
187                                    }
188                                    Arc::new(array_builder.finish())
189                                }
190                                risingwave_pb::plan_common::additional_column::ColumnType::Filename(_) => {
191                                    let mut array_builder = ArrayBuilderImpl::with_type(column_size, source_column.data_type.clone());
192                                    let datum: Datum = Some(ScalarImpl::Utf8(self.file_name.clone().into()));
193                                    array_builder.append_n(record_batch.num_rows(), datum);
194                                    Arc::new(array_builder.finish())
195                                }
196                                _ => unreachable!(),
197                            }
198                        } else {
199                            // For columns defined in the source schema but not present in the Parquet file, null values are filled in.
200                            let mut array_builder =
201                                ArrayBuilderImpl::with_type(column_size, rw_data_type.clone());
202                            array_builder.append_n_null(record_batch.num_rows());
203                            if let Some(metrics) = parquet_source_skip_row_count_metrics.clone() {
204                                metrics.inc_by(record_batch.num_rows() as u64);
205                            }
206                            Arc::new(array_builder.finish())
207                        };
208                        chunk_columns.push(column);
209                    }
210                }
211                crate::source::SourceColumnType::RowId => {
212                    let mut array_builder =
213                        ArrayBuilderImpl::with_type(column_size, source_column.data_type.clone());
214                    let datum: Datum = None;
215                    array_builder.append_n(record_batch.num_rows(), datum);
216                    let res = array_builder.finish();
217                    let column = Arc::new(res);
218                    chunk_columns.push(column);
219                }
220                // The following fields are only used in CDC source
221                crate::source::SourceColumnType::Offset | crate::source::SourceColumnType::Meta => {
222                    unreachable!()
223                }
224            }
225        }
226        if let Some(metrics) = file_source_input_row_count_metrics {
227            metrics.inc_by(record_batch.num_rows() as u64);
228        }
229
230        let data_chunk = DataChunk::new(chunk_columns.clone(), record_batch.num_rows());
231        Ok(data_chunk.into())
232    }
233
234    fn find_parquet_column<'a>(
235        &self,
236        record_batch: &'a RecordBatch,
237        column_name: &str,
238    ) -> Option<(&'a FieldRef, &'a ArrayRef)> {
239        let fields = record_batch.schema_ref().fields();
240        if let Some(index) = fields.iter().position(|field| field.name() == column_name) {
241            return Some((&fields[index], record_batch.column(index)));
242        }
243        if !self.case_insensitive {
244            return None;
245        }
246        let mut matched_index: Option<usize> = None;
247        for (index, field) in fields.iter().enumerate() {
248            if field.name().eq_ignore_ascii_case(column_name) {
249                if matched_index.is_some() {
250                    return None;
251                }
252                matched_index = Some(index);
253            }
254        }
255        matched_index.map(|index| (&fields[index], record_batch.column(index)))
256    }
257}
risingwave_connector/parser/parquet_parser.rs

risingwave_connector/parser/
parquet_parser.rs
