risingwave_connector/source/iceberg/
parquet_file_handler.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::collections::HashMap;
16use std::future::IntoFuture;
17use std::ops::Range;
18use std::pin::Pin;
19use std::sync::Arc;
20
21use anyhow::{Context, bail};
22use bytes::Bytes;
23use futures::future::BoxFuture;
24use futures::{FutureExt, Stream, TryFutureExt};
25use iceberg::io::{
26    FileIOBuilder, FileMetadata, FileRead, S3_ACCESS_KEY_ID, S3_REGION, S3_SECRET_ACCESS_KEY,
27};
28use iceberg_storage_opendal::OpenDalStorageFactory;
29use itertools::Itertools;
30use opendal::Operator;
31use opendal::layers::{LoggingLayer, RetryLayer};
32use opendal::services::{Azblob, Gcs, S3};
33use parquet::arrow::async_reader::AsyncFileReader;
34use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask, parquet_to_arrow_schema};
35use parquet::file::metadata::{FileMetaData, ParquetMetaData, ParquetMetaDataReader};
36use prometheus::core::GenericCounter;
37use risingwave_common::array::StreamChunk;
38use risingwave_common::array::arrow::{IcebergArrowConvert, is_parquet_field_match_source_schema};
39use risingwave_common::catalog::{ColumnDesc, ColumnId};
40use risingwave_common::metrics::LabelGuardedMetric;
41use risingwave_common::util::tokio_util::compat::FuturesAsyncReadCompatExt;
42use url::Url;
43
44use crate::error::ConnectorResult;
45use crate::parser::ParquetParser;
46use crate::source::{Column, SourceColumnDesc};
47
48pub struct ParquetFileReader {
49    meta: FileMetadata,
50    r: Box<dyn FileRead>,
51}
52
53impl ParquetFileReader {
54    pub fn new(meta: FileMetadata, r: Box<dyn FileRead>) -> Self {
55        Self { meta, r }
56    }
57}
58
59impl AsyncFileReader for ParquetFileReader {
60    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, parquet::errors::Result<Bytes>> {
61        Box::pin(
62            self.r
63                .read(range.start as _..range.end as _)
64                .map_err(|err| parquet::errors::ParquetError::External(Box::new(err))),
65        )
66    }
67
68    fn get_metadata(
69        &mut self,
70        _options: Option<&parquet::arrow::arrow_reader::ArrowReaderOptions>,
71    ) -> BoxFuture<'_, parquet::errors::Result<Arc<ParquetMetaData>>> {
72        async move {
73            let reader = ParquetMetaDataReader::new();
74            let size = self.meta.size;
75            let meta = reader.load_and_finish(self, size).await?;
76
77            Ok(Arc::new(meta))
78        }
79        .boxed()
80    }
81}
82
83pub async fn create_parquet_stream_builder(
84    s3_region: String,
85    s3_access_key: String,
86    s3_secret_key: String,
87    location: String,
88) -> Result<ParquetRecordBatchStreamBuilder<ParquetFileReader>, anyhow::Error> {
89    let mut props = HashMap::new();
90    props.insert(S3_REGION, s3_region.clone());
91    props.insert(S3_ACCESS_KEY_ID, s3_access_key.clone());
92    props.insert(S3_SECRET_ACCESS_KEY, s3_secret_key.clone());
93
94    let file_io = FileIOBuilder::new(Arc::new(OpenDalStorageFactory::s3()))
95        .with_props(props)
96        .build();
97    let parquet_file = file_io.new_input(&location)?;
98
99    let parquet_metadata = parquet_file.metadata().await?;
100    let parquet_reader = parquet_file.reader().await?;
101    let parquet_file_reader = ParquetFileReader::new(parquet_metadata, parquet_reader);
102
103    ParquetRecordBatchStreamBuilder::new(parquet_file_reader)
104        .await
105        .map_err(Into::into)
106}
107
108pub fn new_s3_operator(
109    s3_region: String,
110    s3_access_key: String,
111    s3_secret_key: String,
112    bucket: String,
113    s3_endpoint: String,
114) -> ConnectorResult<Operator> {
115    let mut builder = S3::default();
116    builder = builder
117        .region(&s3_region)
118        .endpoint(&s3_endpoint)
119        .access_key_id(&s3_access_key)
120        .secret_access_key(&s3_secret_key)
121        .bucket(&bucket)
122        .disable_config_load();
123    let op: Operator = Operator::new(builder)?
124        .layer(LoggingLayer::default())
125        .layer(RetryLayer::default());
126
127    Ok(op)
128}
129
130pub fn new_gcs_operator(credential: String, bucket: String) -> ConnectorResult<Operator> {
131    // Create gcs builder.
132    let builder = Gcs::default().bucket(&bucket).credential(&credential);
133
134    let operator: Operator = Operator::new(builder)?
135        .layer(LoggingLayer::default())
136        .layer(RetryLayer::default());
137    Ok(operator)
138}
139
140pub fn new_azblob_operator(
141    endpoint: String,
142    account_name: String,
143    account_key: String,
144    container_name: String,
145) -> ConnectorResult<Operator> {
146    // Create azblob builder.
147    let mut builder = Azblob::default();
148    builder = builder
149        .container(&container_name)
150        .endpoint(&endpoint)
151        .account_name(&account_name)
152        .account_key(&account_key);
153
154    let operator: Operator = Operator::new(builder)?
155        .layer(LoggingLayer::default())
156        .layer(RetryLayer::default());
157    Ok(operator)
158}
159
160#[derive(Debug, Clone)]
161pub enum FileScanBackend {
162    S3,
163    Gcs,
164    Azblob,
165}
166
167pub fn extract_bucket_and_file_name(
168    location: &str,
169    file_scan_backend: &FileScanBackend,
170) -> ConnectorResult<(String, String)> {
171    let url = Url::parse(location)?;
172    let bucket = url
173        .host_str()
174        .with_context(|| format!("Invalid url: {}, missing bucket", location))?
175        .to_owned();
176    let prefix = match file_scan_backend {
177        FileScanBackend::S3 => format!("s3://{}/", bucket),
178        FileScanBackend::Gcs => format!("gcs://{}/", bucket),
179        FileScanBackend::Azblob => format!("azblob://{}/", bucket),
180    };
181    let file_name = location[prefix.len()..].to_string();
182    Ok((bucket, file_name))
183}
184
185pub async fn list_data_directory(
186    op: Operator,
187    dir: String,
188    file_scan_backend: &FileScanBackend,
189) -> Result<Vec<String>, anyhow::Error> {
190    let (bucket, file_name) = extract_bucket_and_file_name(&dir, file_scan_backend)?;
191    let prefix = match file_scan_backend {
192        FileScanBackend::S3 => format!("s3://{}/", bucket),
193        FileScanBackend::Gcs => format!("gcs://{}/", bucket),
194        FileScanBackend::Azblob => format!("azblob://{}/", bucket),
195    };
196    if dir.starts_with(&prefix) {
197        op.list(&file_name).await.map_err(Into::into).map(|list| {
198            list.into_iter()
199                .map(|entry| prefix.clone() + entry.path())
200                .collect()
201        })
202    } else {
203        bail!("Invalid url: {}, should start with {}", dir, prefix)
204    }
205}
206
207/// Extracts a suitable `ProjectionMask` from a Parquet file schema based on the user's requested schema.
208///
209/// This function is utilized for column pruning of Parquet files. It checks the user's requested schema
210/// against the schema of the currently read Parquet file. If the provided `columns` are `None`
211/// or if the Parquet file contains nested data types, it returns `ProjectionMask::all()`. Otherwise,
212/// it returns only the columns where both the data type and column name match the requested schema,
213/// facilitating efficient reading of the `RecordBatch`.
214///
215/// # Parameters
216/// - `columns`: An optional vector of `Column` representing the user's requested schema.
217/// - `metadata`: A reference to `FileMetaData` containing the schema and metadata of the Parquet file.
218///
219/// # Returns
220/// - A `ConnectorResult<ProjectionMask>`, which represents the valid columns in the Parquet file schema
221///   that correspond to the requested schema. If an error occurs during processing, it returns an
222///   appropriate error.
223pub fn get_project_mask(
224    columns: Option<Vec<Column>>,
225    metadata: &FileMetaData,
226    case_insensitive: bool,
227) -> ConnectorResult<ProjectionMask> {
228    match columns {
229        Some(rw_columns) => {
230            let root_column_names = metadata
231                .schema_descr()
232                .root_schema()
233                .get_fields()
234                .iter()
235                .map(|field| field.name())
236                .collect_vec();
237
238            let converted_arrow_schema =
239                parquet_to_arrow_schema(metadata.schema_descr(), metadata.key_value_metadata())
240                    .map_err(anyhow::Error::from)?;
241            let mut lowercase_name_to_index: HashMap<String, Option<usize>> = HashMap::new();
242            if case_insensitive {
243                for (index, name) in root_column_names.iter().enumerate() {
244                    let key = name.to_ascii_lowercase();
245                    match lowercase_name_to_index.get(&key) {
246                        Some(Some(_)) => {
247                            lowercase_name_to_index.insert(key, None);
248                        }
249                        Some(None) => {}
250                        None => {
251                            lowercase_name_to_index.insert(key, Some(index));
252                        }
253                    }
254                }
255            }
256
257            let valid_column_indices: Vec<usize> = rw_columns
258                .iter()
259                .filter_map(|column| {
260                    let exact_pos = root_column_names
261                        .iter()
262                        .position(|&name| name == column.name);
263                    let pos = match (case_insensitive, exact_pos) {
264                        (_, Some(pos)) => Some(pos),
265                        (true, None) => lowercase_name_to_index
266                            .get(&column.name.to_ascii_lowercase())
267                            .copied()
268                            .flatten(),
269                        _ => None,
270                    }?;
271                    let arrow_field = converted_arrow_schema.fields.get(pos)?;
272                    let rw_data_type: &risingwave_common::types::DataType = &column.data_type;
273                    if is_parquet_field_match_source_schema(arrow_field, rw_data_type) {
274                        Some(pos)
275                    } else {
276                        // The parquet column exists but its type does not match the declared
277                        // schema; it is NULL-filled rather than decoded to a diverging type.
278                        tracing::warn!(
279                            column = %column.name,
280                            declared_type = %rw_data_type,
281                            parquet_type = %arrow_field.data_type(),
282                            parquet_extension = ?arrow_field.extension_type_name(),
283                            "parquet column does not match the declared source schema; it will be NULL-filled",
284                        );
285                        None
286                    }
287                })
288                .collect();
289
290            Ok(ProjectionMask::roots(
291                metadata.schema_descr(),
292                valid_column_indices,
293            ))
294        }
295        None => Ok(ProjectionMask::all()),
296    }
297}
298
299/// Reads a specified Parquet file and converts its content into a stream of chunks.
300pub async fn read_parquet_file(
301    op: Operator,
302    file_name: String,
303    rw_columns: Option<Vec<Column>>,
304    parser_columns: Option<Vec<SourceColumnDesc>>,
305    case_insensitive: bool,
306    batch_size: usize,
307    offset: usize,
308    file_source_input_row_count_metrics: Option<
309        LabelGuardedMetric<GenericCounter<prometheus::core::AtomicU64>>,
310    >,
311    parquet_source_skip_row_count_metrics: Option<
312        LabelGuardedMetric<GenericCounter<prometheus::core::AtomicU64>>,
313    >,
314) -> ConnectorResult<
315    Pin<Box<dyn Stream<Item = Result<StreamChunk, crate::error::ConnectorError>> + Send>>,
316> {
317    let mut reader: tokio_util::compat::Compat<opendal::FuturesAsyncReader> = op
318        .reader_with(&file_name)
319        .into_future() // Unlike `rustc`, `try_stream` seems require manual `into_future`.
320        .await?
321        .into_futures_async_read(..)
322        .await?
323        .compat();
324    let parquet_metadata = reader
325        .get_metadata(None)
326        .await
327        .map_err(anyhow::Error::from)?;
328
329    let file_metadata = parquet_metadata.file_metadata();
330    {
331        // Log parquet file-level metadata
332        tracing::info!(
333            "Reading parquet file: {}, from offset {}, num_row_groups={}, total_rows={}, kv_len={}",
334            file_name,
335            offset,
336            parquet_metadata.row_groups().len(),
337            file_metadata.num_rows(),
338            file_metadata
339                .key_value_metadata()
340                .map(|m| m.len())
341                .unwrap_or(0)
342        );
343        // Log each leaf column's path and types
344        let schema_descr = file_metadata.schema_descr();
345        for col in schema_descr.columns() {
346            let path = col.path().string();
347            let physical = col.physical_type();
348            let logical = col.logical_type_ref();
349            tracing::debug!(
350                file = %file_name,
351                column_path = path,
352                physical_type = ?physical,
353                logical_type = ?logical,
354                type_length = ?col.type_length(),
355                max_def_level = col.max_def_level(),
356                max_rep_level = col.max_rep_level(),
357                "Parquet file column schema: "
358            );
359        }
360    }
361    let projection_mask = get_project_mask(rw_columns, file_metadata, case_insensitive)?;
362
363    // For the Parquet format, we directly convert from a record batch to a stream chunk.
364    // Therefore, the offset of the Parquet file represents the current position in terms of the number of rows read from the file.
365    let record_batch_stream = ParquetRecordBatchStreamBuilder::new(reader)
366        .await?
367        .with_batch_size(batch_size)
368        .with_projection(projection_mask)
369        .with_offset(offset)
370        .build()?;
371    let converted_arrow_schema = parquet_to_arrow_schema(
372        file_metadata.schema_descr(),
373        file_metadata.key_value_metadata(),
374    )?;
375    let columns = match parser_columns {
376        Some(columns) => columns,
377        None => converted_arrow_schema
378            .fields
379            .iter()
380            .enumerate()
381            .map(|(index, field_ref)| {
382                let data_type = IcebergArrowConvert
383                    .type_from_field(field_ref)
384                    .with_context(|| {
385                        format!(
386                            "cannot infer the RisingWave type of parquet column `{}` in file {}",
387                            field_ref.name(),
388                            file_name
389                        )
390                    })?;
391                let column_desc = ColumnDesc::named(
392                    field_ref.name().clone(),
393                    ColumnId::new(index as i32),
394                    data_type,
395                );
396                Ok(SourceColumnDesc::from(&column_desc))
397            })
398            .collect::<ConnectorResult<Vec<_>>>()?,
399    };
400    let parquet_parser = ParquetParser::new(columns, file_name, offset, case_insensitive)?;
401    let msg_stream: Pin<
402        Box<dyn Stream<Item = Result<StreamChunk, crate::error::ConnectorError>> + Send>,
403    > = parquet_parser.into_stream(
404        record_batch_stream,
405        file_source_input_row_count_metrics,
406        parquet_source_skip_row_count_metrics,
407    );
408    Ok(msg_stream)
409}
410
411pub async fn get_parquet_fields(
412    op: Operator,
413    file_name: String,
414) -> ConnectorResult<risingwave_common::array::arrow::arrow_schema_iceberg::Fields> {
415    let mut reader: tokio_util::compat::Compat<opendal::FuturesAsyncReader> = op
416        .reader_with(&file_name)
417        .into_future() // Unlike `rustc`, `try_stream` seems require manual `into_future`.
418        .await?
419        .into_futures_async_read(..)
420        .await?
421        .compat();
422    let parquet_metadata = reader.get_metadata(None).await?;
423
424    let file_metadata = parquet_metadata.file_metadata();
425    let converted_arrow_schema = parquet_to_arrow_schema(
426        file_metadata.schema_descr(),
427        file_metadata.key_value_metadata(),
428    )?;
429    let fields: risingwave_common::array::arrow::arrow_schema_iceberg::Fields =
430        converted_arrow_schema.fields;
431    Ok(fields)
432}
risingwave_connector/source/iceberg/parquet_file_handler.rs

risingwave_connector/source/iceberg/
parquet_file_handler.rs
