Repository: risingwavelabs/risingwave (language: python)

Issue:
feat(connector): support more parquet data types
I hereby agree to the terms of the [RisingWave Labs, Inc. Contributor License Agreement](https://gist.github.com/TennyZhuang/f00be7f16996ea48effb049aa7be4d66#file-rw_cla).

## What's changed and what's your intention?
Since RisingWave does not have Uint data type, previously we do not support `UIntx`(`UInt8, UInt16, UInt32, UInt64`) in parquet encode, and `Float16` is also not supported.

However, user have a bad experience when they encounter an error of unsupported data type when using file_scan() or parquet file source. 

This pr add some types that can be workarounded:
```
Int8  -> Int16
UInt8  -> Int16
UInt16 -> Int32
UInt32 -> Int64
UInt64 -> Decimal
Float16 -> Float32
```
For example, if the Parquet file that the user needs to read contains a column of Uint64 type, this column should be defined as Decimal type when creating the RisingWave table.


<!--

**Please do not leave this empty!**

Please explain **IN DETAIL** what the changes are in this PR and why they are needed:

- Summarize your change (**mandatory**)
- How does this PR work? Need a brief introduction for the changed logic (optional)
- Describe clearly one logical change and avoid lazy messages (optional)
- Describe any limitations of the current code (optional)
- Refer to a related PR or issue link (optional)

-->

## Checklist

- [ ] I have written necessary rustdoc comments
- [ ] I have added necessary unit tests and integration tests
- [ ] I have added test labels as necessary. See [details](https://github.com/risingwavelabs/risingwave/blob/main/docs/developer-guide.md#ci-labels-guide).
- [ ] I have added fuzzing tests or opened an issue to track them. (Optional, recommended for new SQL features #7934).
- [ ] My PR contains breaking changes. (If it deprecates some features, please create a tracking issue to remove them in the future).
- [ ] All checks passed in `./risedev check` (or alias, `./risedev c`)
- [ ] My PR changes performance-critical code. (Please run macro/micro-benchmarks and show the results.)
<!-- To manually trigger a benchmark, please check out [Notion](https://www.notion.so/risingwave-labs/Manually-trigger-nexmark-performance-dashboard-test-b784f1eae1cf48889b2645d020b6b7d3). -->
- [ ] My PR contains critical fixes that are necessary to be merged into the latest release. (Please check out the [details](https://github.com/risingwavelabs/risingwave/blob/main/CONTRIBUTING.md))

## Documentation

- [x] My PR needs documentation updates. (Please use the **Release note** section below to summarize the impact on users)

## Release note

We need the documentation team to note these special conversions and inform us of the list of data types that we do not support.

Perhaps in https://docs.risingwave.com/integrations/sources/apache-iceberg#data-type-mapping


<!--
Please create a release note for your changes.

Discuss technical details in the "What's changed" section, and
focus on the impact on users in the release note.

You should also mention the environment or conditions where the impact may occur.
-->

</details>

Sqlsmith: Sql feature generation
As we add more features, these should be tested via sqlsmith as well.

SQL Features:
- [ ] Generate other table functions: `generate_series`, `unnest`. See: https://github.com/singularity-data/risingwave/issues/4540
- [ ] Generate array_agg, jsonb_agg, jsonb_object_agg
- [ ] Generate select from System Tables
- [ ] Generate `NATURAL JOIN`
- [ ] Generate `JOIN` with `USING`
- [ ] Generate query w/o `GROUP BY`
- [ ] Randomly toggle session variables
- [ ] Generate `pow`/`exp`.
- [ ] Generate TopN by group: https://www.risingwave.dev/docs/current/sql-pattern-topn/
- [ ] Generate jsonb exprs, see: https://github.com/risingwavelabs/risingwave/issues/7714
- [ ] #7132
- [x] Generate order-by-clauses with `NULLS { FIRST | LAST }` (will be done in https://github.com/risingwavelabs/risingwave/pull/10864/commits/000cbc11cbc04f20abfd127b1bf4b3fd24e75c44)
- [ ] Generate `max`,`min`,`count` for `timestamptz`
- [ ] Generate `bit_and`, `bit_or` and `big_xor` in batch queries, and `bit_xor` in streaming queries.
- [ ] Generate array range access.
- [ ] https://github.com/risingwavelabs/risingwave/issues/10241
- [ ] Generate window functions other than `row_number`, `rank`.
- [ ] Generate `ALTER` DDL statements.\
- [ ] https://github.com/risingwavelabs/risingwave/issues/10516
- [ ] Generate temporal filter in sqlsmith


Discussion / hints:

Fix patch:
diff --git a/src/batch/src/executor/s3_file_scan.rs b/src/batch/src/executor/s3_file_scan.rs
index 5096d1d625fa..38907c63f841 100644
--- a/src/batch/src/executor/s3_file_scan.rs
+++ b/src/batch/src/executor/s3_file_scan.rs
@@ -12,13 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use anyhow::anyhow;
 use futures_async_stream::try_stream;
 use futures_util::stream::StreamExt;
-use parquet::arrow::ProjectionMask;
-use risingwave_common::array::arrow::IcebergArrowConvert;
 use risingwave_common::catalog::{Field, Schema};
-use risingwave_connector::source::iceberg::parquet_file_reader::create_parquet_stream_builder;
+use risingwave_connector::source::iceberg::{new_s3_operator, read_parquet_file};
 use risingwave_pb::batch_plan::file_scan_node;
 use risingwave_pb::batch_plan::file_scan_node::StorageType;
 use risingwave_pb::batch_plan::plan_node::NodeBody;
@@ -85,34 +82,18 @@ impl S3FileScanExecutor {
     async fn do_execute(self: Box<Self>) {
         assert_eq!(self.file_format, FileFormat::Parquet);
         for file in self.file_location {
-            let mut batch_stream_builder = create_parquet_stream_builder(
+            let op = new_s3_operator(
                 self.s3_region.clone(),
                 self.s3_access_key.clone(),
                 self.s3_secret_key.clone(),
-                file,
-            )
-            .await?;
-
-            let arrow_schema = batch_stream_builder.schema();
-            assert_eq!(arrow_schema.fields.len(), self.schema.fields.len());
-            for (field, arrow_field) in self.schema.fields.iter().zip(arrow_schema.fields.iter()) {
-                assert_eq!(*field.name, *arrow_field.name());
-            }
-
-            batch_stream_builder = batch_stream_builder.with_projection(ProjectionMask::all());
-
-            batch_stream_builder = batch_stream_builder.with_batch_size(self.batch_size);
-
-            let record_batch_stream = batch_stream_builder
-                .build()
-                .map_err(|e| anyhow!(e).context("fail to build arrow stream builder"))?;
-
+                file.clone(),
+            )?;
+            let chunk_stream = read_parquet_file(op, file, None, None, self.batch_size, 0).await?;
             #[for_await]
-            for record_batch in record_batch_stream {
-                let record_batch = record_batch.map_err(BatchError::Parquet)?;
-                let chunk = IcebergArrowConvert.chunk_from_record_batch(&record_batch)?;
-                debug_assert_eq!(chunk.data_types(), self.schema.data_types());
-                yield chunk;
+            for stream_chunk in chunk_stream {
+                let stream_chunk = stream_chunk?;
+                let (data_chunk, _) = stream_chunk.into_parts();
+                yield data_chunk;
             }
         }
     }
diff --git a/src/common/src/array/arrow/arrow_impl.rs b/src/common/src/array/arrow/arrow_impl.rs
index 2e35a1403781..f72a1d8739be 100644
--- a/src/common/src/array/arrow/arrow_impl.rs
+++ b/src/common/src/array/arrow/arrow_impl.rs
@@ -505,6 +505,12 @@ pub trait FromArrow {
             Int16 => DataType::Int16,
             Int32 => DataType::Int32,
             Int64 => DataType::Int64,
+            Int8 => DataType::Int16,
+            UInt8 => DataType::Int16,
+            UInt16 => DataType::Int32,
+            UInt32 => DataType::Int64,
+            UInt64 => DataType::Decimal,
+            Float16 => DataType::Float32,
             Float32 => DataType::Float32,
             Float64 => DataType::Float64,
             Decimal128(_, _) => DataType::Decimal,
@@ -581,11 +587,18 @@ pub trait FromArrow {
         }
         match array.data_type() {
             Boolean => self.from_bool_array(array.as_any().downcast_ref().unwrap()),
+            Int8 => self.from_int8_array(array.as_any().downcast_ref().unwrap()),
             Int16 => self.from_int16_array(array.as_any().downcast_ref().unwrap()),
             Int32 => self.from_int32_array(array.as_any().downcast_ref().unwrap()),
             Int64 => self.from_int64_array(array.as_any().downcast_ref().unwrap()),
+            UInt8 => self.from_uint8_array(array.as_any().downcast_ref().unwrap()),
+            UInt16 => self.from_uint16_array(array.as_any().downcast_ref().unwrap()),
+            UInt32 => self.from_uint32_array(array.as_any().downcast_ref().unwrap()),
+
+            UInt64 => self.from_uint64_array(array.as_any().downcast_ref().unwrap()),
             Decimal128(_, _) => self.from_decimal128_array(array.as_any().downcast_ref().unwrap()),
             Decimal256(_, _) => self.from_int256_array(array.as_any().downcast_ref().unwrap()),
+            Float16 => self.from_float16_array(array.as_any().downcast_ref().unwrap()),
             Float32 => self.from_float32_array(array.as_any().downcast_ref().unwrap()),
             Float64 => self.from_float64_array(array.as_any().downcast_ref().unwrap()),
             Date32 => self.from_date32_array(array.as_any().downcast_ref().unwrap()),
@@ -669,6 +682,22 @@ pub trait FromArrow {
         Ok(ArrayImpl::Int16(array.into()))
     }
 
+    fn from_int8_array(&self, array: &arrow_array::Int8Array) -> Result<ArrayImpl, ArrayError> {
+        Ok(ArrayImpl::Int16(array.into()))
+    }
+
+    fn from_uint8_array(&self, array: &arrow_array::UInt8Array) -> Result<ArrayImpl, ArrayError> {
+        Ok(ArrayImpl::Int16(array.into()))
+    }
+
+    fn from_uint16_array(&self, array: &arrow_array::UInt16Array) -> Result<ArrayImpl, ArrayError> {
+        Ok(ArrayImpl::Int32(array.into()))
+    }
+
+    fn from_uint32_array(&self, array: &arrow_array::UInt32Array) -> Result<ArrayImpl, ArrayError> {
+        Ok(ArrayImpl::Int64(array.into()))
+    }
+
     fn from_int32_array(&self, array: &arrow_array::Int32Array) -> Result<ArrayImpl, ArrayError> {
         Ok(ArrayImpl::Int32(array.into()))
     }
@@ -691,6 +720,17 @@ pub trait FromArrow {
         Ok(ArrayImpl::Decimal(array.try_into()?))
     }
 
+    fn from_uint64_array(&self, array: &arrow_array::UInt64Array) -> Result<ArrayImpl, ArrayError> {
+        Ok(ArrayImpl::Decimal(array.try_into()?))
+    }
+
+    fn from_float16_array(
+        &self,
+        array: &arrow_array::Float16Array,
+    ) -> Result<ArrayImpl, ArrayError> {
+        Ok(ArrayImpl::Float32(array.try_into()?))
+    }
+
     fn from_float32_array(
         &self,
         array: &arrow_array::Float32Array,
@@ -908,6 +948,37 @@ macro_rules! converts {
     };
 }
 
+/// Used to convert different types.
+macro_rules! converts_with_type {
+    ($ArrayType:ty, $ArrowType:ty, $FromType:ty, $ToType:ty) => {
+        impl From<&$ArrayType> for $ArrowType {
+            fn from(array: &$ArrayType) -> Self {
+                let values: Vec<Option<$ToType>> =
+                    array.iter().map(|x| x.map(|v| v as $ToType)).collect();
+                <$ArrowType>::from_iter(values)
+            }
+        }
+
+        impl From<&$ArrowType> for $ArrayType {
+            fn from(array: &$ArrowType) -> Self {
+                let values: Vec<Option<$FromType>> =
+                    array.iter().map(|x| x.map(|v| v as $FromType)).collect();
+                <$ArrayType>::from_iter(values)
+            }
+        }
+
+        impl From<&[$ArrowType]> for $ArrayType {
+            fn from(arrays: &[$ArrowType]) -> Self {
+                let values: Vec<Option<$FromType>> = arrays
+                    .iter()
+                    .flat_map(|a| a.iter().map(|x| x.map(|v| v as $FromType)))
+                    .collect();
+                <$ArrayType>::from_iter(values)
+            }
+        }
+    };
+}
+
 macro_rules! converts_with_timeunit {
     ($ArrayType:ty, $ArrowType:ty, $time_unit:expr, @map) => {
 
@@ -960,6 +1031,11 @@ converts!(TimeArray, arrow_array::Time64MicrosecondArray, @map);
 converts!(IntervalArray, arrow_array::IntervalMonthDayNanoArray, @map);
 converts!(SerialArray, arrow_array::Int64Array, @map);
 
+converts_with_type!(I16Array, arrow_array::Int8Array, i16, i8);
+converts_with_type!(I16Array, arrow_array::UInt8Array, i16, u8);
+converts_with_type!(I32Array, arrow_array::UInt16Array, i32, u16);
+converts_with_type!(I64Array, arrow_array::UInt32Array, i64, u32);
+
 converts_with_timeunit!(TimestampArray, arrow_array::TimestampSecondArray, TimeUnit::Second, @map);
 converts_with_timeunit!(TimestampArray, arrow_array::TimestampMillisecondArray, TimeUnit::Millisecond, @map);
 converts_with_timeunit!(TimestampArray, arrow_array::TimestampMicrosecondArray, TimeUnit::Microsecond, @map);
@@ -1171,6 +1247,38 @@ impl TryFrom<&arrow_array::Decimal128Array> for DecimalArray {
     }
 }
 
+// Since RisingWave does not support UInt type, convert UInt64Array to Decimal.
+impl TryFrom<&arrow_array::UInt64Array> for DecimalArray {
+    type Error = ArrayError;
+
+    fn try_from(array: &arrow_array::UInt64Array) -> Result<Self, Self::Error> {
+        let from_arrow = |value| {
+            // Convert the value to a Decimal with scale 0
+            let res = Decimal::from(value);
+            Ok(res)
+        };
+
+        // Map over the array and convert each value
+        array
+            .iter()
+            .map(|o| o.map(from_arrow).transpose())
+            .collect::<Result<Self, Self::Error>>()
+    }
+}
+
+impl TryFrom<&arrow_array::Float16Array> for F32Array {
+    type Error = ArrayError;
+
+    fn try_from(array: &arrow_array::Float16Array) -> Result<Self, Self::Error> {
+        let from_arrow = |value| Ok(f32::from(value));
+
+        array
+            .iter()
+            .map(|o| o.map(from_arrow).transpose())
+            .collect::<Result<Self, Self::Error>>()
+    }
+}
+
 impl TryFrom<&arrow_array::LargeBinaryArray> for DecimalArray {
     type Error = ArrayError;
 
@@ -1311,6 +1419,7 @@ impl From<&arrow_array::Decimal256Array> for Int256Array {
 
 #[cfg(test)]
 mod tests {
+
     use super::*;
 
     #[test]
@@ -1355,6 +1464,50 @@ mod tests {
         assert_eq!(F64Array::from(&arrow), array);
     }
 
+    #[test]
+    fn int8() {
+        let array: PrimitiveArray<i16> = I16Array::from_iter([None, Some(-128), Some(127)]);
+        let arr = arrow_array::Int8Array::from(vec![None, Some(-128), Some(127)]);
+        let converted: PrimitiveArray<i16> = (&arr).into();
+        assert_eq!(converted, array);
+    }
+
+    #[test]
+    fn uint8() {
+        let array: PrimitiveArray<i16> = I16Array::from_iter([None, Some(7), Some(25)]);
+        let arr = arrow_array::UInt8Array::from(vec![None, Some(7), Some(25)]);
+        let converted: PrimitiveArray<i16> = (&arr).into();
+        assert_eq!(converted, array);
+    }
+
+    #[test]
+    fn uint16() {
+        let array: PrimitiveArray<i32> = I32Array::from_iter([None, Some(7), Some(65535)]);
+        let arr = arrow_array::UInt16Array::from(vec![None, Some(7), Some(65535)]);
+        let converted: PrimitiveArray<i32> = (&arr).into();
+        assert_eq!(converted, array);
+    }
+
+    #[test]
+    fn uint32() {
+        let array: PrimitiveArray<i64> = I64Array::from_iter([None, Some(7), Some(4294967295)]);
+        let arr = arrow_array::UInt32Array::from(vec![None, Some(7), Some(4294967295)]);
+        let converted: PrimitiveArray<i64> = (&arr).into();
+        assert_eq!(converted, array);
+    }
+
+    #[test]
+    fn uint64() {
+        let array: PrimitiveArray<Decimal> = DecimalArray::from_iter([
+            None,
+            Some(Decimal::Normalized("7".parse().unwrap())),
+            Some(Decimal::Normalized("18446744073709551615".parse().unwrap())),
+        ]);
+        let arr = arrow_array::UInt64Array::from(vec![None, Some(7), Some(18446744073709551615)]);
+        let converted: PrimitiveArray<Decimal> = (&arr).try_into().unwrap();
+        assert_eq!(converted, array);
+    }
+
     #[test]
     fn date() {
         let array = DateArray::from_iter([
diff --git a/src/connector/src/source/filesystem/opendal_source/opendal_reader.rs b/src/connector/src/source/filesystem/opendal_source/opendal_reader.rs
index 14258d892465..ca8ee1ae486b 100644
--- a/src/connector/src/source/filesystem/opendal_source/opendal_reader.rs
+++ b/src/connector/src/source/filesystem/opendal_source/opendal_reader.rs
@@ -19,24 +19,19 @@ use async_compression::tokio::bufread::GzipDecoder;
 use async_trait::async_trait;
 use futures::TryStreamExt;
 use futures_async_stream::try_stream;
-use itertools::Itertools;
 use opendal::Operator;
-use parquet::arrow::async_reader::AsyncFileReader;
-use parquet::arrow::{parquet_to_arrow_schema, ParquetRecordBatchStreamBuilder, ProjectionMask};
-use parquet::file::metadata::FileMetaData;
-use risingwave_common::array::arrow::IcebergArrowConvert;
 use risingwave_common::array::StreamChunk;
-use risingwave_common::util::tokio_util::compat::FuturesAsyncReadCompatExt;
 use tokio::io::{AsyncRead, BufReader};
 use tokio_util::io::{ReaderStream, StreamReader};
 
 use super::opendal_enumerator::OpendalEnumerator;
 use super::OpendalSource;
 use crate::error::ConnectorResult;
-use crate::parser::{ByteStreamSourceParserImpl, EncodingProperties, ParquetParser, ParserConfig};
+use crate::parser::{ByteStreamSourceParserImpl, EncodingProperties, ParserConfig};
 use crate::source::filesystem::file_common::CompressionFormat;
 use crate::source::filesystem::nd_streaming::need_nd_streaming;
 use crate::source::filesystem::{nd_streaming, OpendalFsSplit};
+use crate::source::iceberg::read_parquet_file;
 use crate::source::{
     BoxChunkSourceStream, Column, SourceContextRef, SourceMessage, SourceMeta, SplitMetaData,
     SplitReader,
@@ -91,38 +86,15 @@ impl<Src: OpendalSource> OpendalReader<Src> {
             let msg_stream;
 
             if let EncodingProperties::Parquet = &self.parser_config.specific.encoding_config {
-                // // If the format is "parquet", use `ParquetParser` to convert `record_batch` into stream chunk.
-                let mut reader: tokio_util::compat::Compat<opendal::FuturesAsyncReader> = self
-                    .connector
-                    .op
-                    .reader_with(&object_name)
-                    .into_future() // Unlike `rustc`, `try_stream` seems require manual `into_future`.
-                    .await?
-                    .into_futures_async_read(..)
-                    .await?
-                    .compat();
-                let parquet_metadata = reader.get_metadata().await.map_err(anyhow::Error::from)?;
-
-                let file_metadata = parquet_metadata.file_metadata();
-                let column_indices =
-                    extract_valid_column_indices(self.columns.clone(), file_metadata)?;
-                let projection_mask =
-                    ProjectionMask::leaves(file_metadata.schema_descr(), column_indices);
-                // For the Parquet format, we directly convert from a record batch to a stream chunk.
-                // Therefore, the offset of the Parquet file represents the current position in terms of the number of rows read from the file.
-                let record_batch_stream = ParquetRecordBatchStreamBuilder::new(reader)
-                    .await?
-                    .with_batch_size(self.source_ctx.source_ctrl_opts.chunk_size)
-                    .with_projection(projection_mask)
-                    .with_offset(split.offset)
-                    .build()?;
-
-                let parquet_parser = ParquetParser::new(
-                    self.parser_config.common.rw_columns.clone(),
+                msg_stream = read_parquet_file(
+                    self.connector.op.clone(),
                     object_name,
+                    self.columns.clone(),
+                    Some(self.parser_config.common.rw_columns.clone()),
+                    self.source_ctx.source_ctrl_opts.chunk_size,
                     split.offset,
-                )?;
-                msg_stream = parquet_parser.into_stream(record_batch_stream);
+                )
+                .await?;
             } else {
                 let data_stream = Self::stream_read_object(
                     self.connector.op.clone(),
@@ -229,61 +201,3 @@ impl<Src: OpendalSource> OpendalReader<Src> {
         }
     }
 }
-
-/// Extracts valid column indices from a Parquet file schema based on the user's requested schema.
-///
-/// This function is used for column pruning of Parquet files. It calculates the intersection
-/// between the columns in the currently read Parquet file and the schema provided by the user.
-/// This is useful for reading a `RecordBatch` with the appropriate `ProjectionMask`, ensuring that
-/// only the necessary columns are read.
-///
-/// # Parameters
-/// - `columns`: A vector of `Column` representing the user's requested schema.
-/// - `metadata`: A reference to `FileMetaData` containing the schema and metadata of the Parquet file.
-///
-/// # Returns
-/// - A `ConnectorResult<Vec<usize>>`, which contains the indices of the valid columns in the
-///   Parquet file schema that match the requested schema. If an error occurs during processing,
-///   it returns an appropriate error.
-pub fn extract_valid_column_indices(
-    columns: Option<Vec<Column>>,
-    metadata: &FileMetaData,
-) -> ConnectorResult<Vec<usize>> {
-    match columns {
-        Some(rw_columns) => {
-            let parquet_column_names = metadata
-                .schema_descr()
-                .columns()
-                .iter()
-                .map(|c| c.name())
-                .collect_vec();
-
-            let converted_arrow_schema =
-                parquet_to_arrow_schema(metadata.schema_descr(), metadata.key_value_metadata())
-                    .map_err(anyhow::Error::from)?;
-
-            let valid_column_indices: Vec<usize> = rw_columns
-                .iter()
-                .filter_map(|column| {
-                    parquet_column_names
-                        .iter()
-                        .position(|&name| name == column.name)
-                        .and_then(|pos| {
-                            // We should convert Arrow field to the rw data type instead of converting the rw data type to the Arrow data type for comparison.
-                            // The reason is that for the timestamp type, the different time units in Arrow need to match with the timestamp and timestamptz in rw.
-                            let arrow_filed_to_rw_data_type = IcebergArrowConvert
-                                .type_from_field(converted_arrow_schema.field(pos))
-                                .ok()?;
-                            if arrow_filed_to_rw_data_type == column.data_type {
-                                Some(pos)
-                            } else {
-                                None
-                            }
-                        })
-                })
-                .collect();
-            Ok(valid_column_indices)
-        }
-        None => Ok(vec![]),
-    }
-}
diff --git a/src/connector/src/source/iceberg/mod.rs b/src/connector/src/source/iceberg/mod.rs
index 60a26e43e1d3..ad37120170a8 100644
--- a/src/connector/src/source/iceberg/mod.rs
+++ b/src/connector/src/source/iceberg/mod.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub mod parquet_file_reader;
+pub mod parquet_file_handler;
 
 use std::collections::HashMap;
 
@@ -23,7 +23,7 @@ use iceberg::scan::FileScanTask;
 use iceberg::spec::TableMetadata;
 use iceberg::table::Table;
 use itertools::Itertools;
-pub use parquet_file_reader::*;
+pub use parquet_file_handler::*;
 use risingwave_common::bail;
 use risingwave_common::catalog::Schema;
 use risingwave_common::types::JsonbVal;
diff --git a/src/connector/src/source/iceberg/parquet_file_handler.rs b/src/connector/src/source/iceberg/parquet_file_handler.rs
new file mode 100644
index 000000000000..c4bfceb1ed89
--- /dev/null
+++ b/src/connector/src/source/iceberg/parquet_file_handler.rs
@@ -0,0 +1,389 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+use std::future::IntoFuture;
+use std::ops::Range;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use anyhow::anyhow;
+use bytes::Bytes;
+use futures::future::BoxFuture;
+use futures::{FutureExt, Stream, TryFutureExt};
+use iceberg::io::{
+    FileIOBuilder, FileMetadata, FileRead, S3_ACCESS_KEY_ID, S3_REGION, S3_SECRET_ACCESS_KEY,
+};
+use iceberg::{Error, ErrorKind};
+use itertools::Itertools;
+use opendal::layers::{LoggingLayer, RetryLayer};
+use opendal::services::S3;
+use opendal::Operator;
+use parquet::arrow::async_reader::AsyncFileReader;
+use parquet::arrow::{parquet_to_arrow_schema, ParquetRecordBatchStreamBuilder, ProjectionMask};
+use parquet::file::metadata::{FileMetaData, ParquetMetaData, ParquetMetaDataReader};
+use risingwave_common::array::arrow::arrow_schema_udf::{DataType as ArrowDateType, IntervalUnit};
+use risingwave_common::array::arrow::IcebergArrowConvert;
+use risingwave_common::array::StreamChunk;
+use risingwave_common::catalog::ColumnId;
+use risingwave_common::types::DataType as RwDataType;
+use risingwave_common::util::tokio_util::compat::FuturesAsyncReadCompatExt;
+use url::Url;
+
+use crate::error::ConnectorResult;
+use crate::parser::ParquetParser;
+use crate::source::{Column, SourceColumnDesc};
+
+pub struct ParquetFileReader<R: FileRead> {
+    meta: FileMetadata,
+    r: R,
+}
+
+impl<R: FileRead> ParquetFileReader<R> {
+    pub fn new(meta: FileMetadata, r: R) -> Self {
+        Self { meta, r }
+    }
+}
+
+impl<R: FileRead> AsyncFileReader for ParquetFileReader<R> {
+    fn get_bytes(&mut self, range: Range<usize>) -> BoxFuture<'_, parquet::errors::Result<Bytes>> {
+        Box::pin(
+            self.r
+                .read(range.start as _..range.end as _)
+                .map_err(|err| parquet::errors::ParquetError::External(Box::new(err))),
+        )
+    }
+
+    fn get_metadata(&mut self) -> BoxFuture<'_, parquet::errors::Result<Arc<ParquetMetaData>>> {
+        async move {
+            let reader = ParquetMetaDataReader::new();
+            let size = self.meta.size as usize;
+            let meta = reader.load_and_finish(self, size).await?;
+
+            Ok(Arc::new(meta))
+        }
+        .boxed()
+    }
+}
+
+pub async fn create_parquet_stream_builder(
+    s3_region: String,
+    s3_access_key: String,
+    s3_secret_key: String,
+    location: String,
+) -> Result<ParquetRecordBatchStreamBuilder<ParquetFileReader<impl FileRead>>, anyhow::Error> {
+    let mut props = HashMap::new();
+    props.insert(S3_REGION, s3_region.clone());
+    props.insert(S3_ACCESS_KEY_ID, s3_access_key.clone());
+    props.insert(S3_SECRET_ACCESS_KEY, s3_secret_key.clone());
+
+    let file_io_builder = FileIOBuilder::new("s3");
+    let file_io = file_io_builder
+        .with_props(props.into_iter())
+        .build()
+        .map_err(|e| anyhow!(e))?;
+    let parquet_file = file_io.new_input(&location).map_err(|e| anyhow!(e))?;
+
+    let parquet_metadata = parquet_file.metadata().await.map_err(|e| anyhow!(e))?;
+    let parquet_reader = parquet_file.reader().await.map_err(|e| anyhow!(e))?;
+    let parquet_file_reader = ParquetFileReader::new(parquet_metadata, parquet_reader);
+
+    ParquetRecordBatchStreamBuilder::new(parquet_file_reader)
+        .await
+        .map_err(|e| anyhow!(e))
+}
+
+pub fn new_s3_operator(
+    s3_region: String,
+    s3_access_key: String,
+    s3_secret_key: String,
+    location: String,
+) -> ConnectorResult<Operator> {
+    // Create s3 builder.
+    let bucket = extract_bucket(&location);
+    let mut builder = S3::default().bucket(&bucket).region(&s3_region);
+    builder = builder.secret_access_key(&s3_access_key);
+    builder = builder.secret_access_key(&s3_secret_key);
+    builder = builder.endpoint(&format!(
+        "https://{}.s3.{}.amazonaws.com",
+        bucket, s3_region
+    ));
+
+    builder = builder.disable_config_load();
+
+    let op: Operator = Operator::new(builder)?
+        .layer(LoggingLayer::default())
+        .layer(RetryLayer::default())
+        .finish();
+
+    Ok(op)
+}
+
+fn extract_bucket(location: &str) -> String {
+    let prefix = "s3://";
+    let start = prefix.len();
+    let end = location[start..]
+        .find('/')
+        .unwrap_or(location.len() - start);
+    location[start..start + end].to_string()
+}
+
+pub async fn list_s3_directory(
+    s3_region: String,
+    s3_access_key: String,
+    s3_secret_key: String,
+    dir: String,
+) -> Result<Vec<String>, anyhow::Error> {
+    let url = Url::parse(&dir)?;
+    let bucket = url.host_str().ok_or_else(|| {
+        Error::new(
+            ErrorKind::DataInvalid,
+            format!("Invalid s3 url: {}, missing bucket", dir),
+        )
+    })?;
+
+    let prefix = format!("s3://{}/", bucket);
+    if dir.starts_with(&prefix) {
+        let mut builder = S3::default();
+        builder = builder
+            .region(&s3_region)
+            .access_key_id(&s3_access_key)
+            .secret_access_key(&s3_secret_key)
+            .bucket(bucket);
+        let op = Operator::new(builder)?
+            .layer(RetryLayer::default())
+            .finish();
+
+        op.list(&dir[prefix.len()..])
+            .await
+            .map_err(|e| anyhow!(e))
+            .map(|list| {
+                list.into_iter()
+                    .map(|entry| prefix.to_string() + entry.path())
+                    .collect()
+            })
+    } else {
+        Err(Error::new(
+            ErrorKind::DataInvalid,
+            format!("Invalid s3 url: {}, should start with {}", dir, prefix),
+        ))?
+    }
+}
+
+/// Extracts valid column indices from a Parquet file schema based on the user's requested schema.
+///
+/// This function is used for column pruning of Parquet files. It calculates the intersection
+/// between the columns in the currently read Parquet file and the schema provided by the user.
+/// This is useful for reading a `RecordBatch` with the appropriate `ProjectionMask`, ensuring that
+/// only the necessary columns are read.
+///
+/// # Parameters
+/// - `columns`: A vector of `Column` representing the user's requested schema.
+/// - `metadata`: A reference to `FileMetaData` containing the schema and metadata of the Parquet file.
+///
+/// # Returns
+/// - A `ConnectorResult<Vec<usize>>`, which contains the indices of the valid columns in the
+///   Parquet file schema that match the requested schema. If an error occurs during processing,
+///   it returns an appropriate error.
+pub fn extract_valid_column_indices(
+    columns: Option<Vec<Column>>,
+    metadata: &FileMetaData,
+) -> ConnectorResult<Vec<usize>> {
+    match columns {
+        Some(rw_columns) => {
+            let parquet_column_names = metadata
+                .schema_descr()
+                .columns()
+                .iter()
+                .map(|c| c.name())
+                .collect_vec();
+
+            let converted_arrow_schema =
+                parquet_to_arrow_schema(metadata.schema_descr(), metadata.key_value_metadata())
+                    .map_err(anyhow::Error::from)?;
+
+            let valid_column_indices: Vec<usize> = rw_columns
+                .iter()
+                .filter_map(|column| {
+                    parquet_column_names
+                        .iter()
+                        .position(|&name| name == column.name)
+                        .and_then(|pos| {
+                            let arrow_data_type: &risingwave_common::array::arrow::arrow_schema_udf::DataType = converted_arrow_schema.field(pos).data_type();
+                            let rw_data_type: &risingwave_common::types::DataType = &column.data_type;
+
+                            if is_parquet_schema_match_source_schema(arrow_data_type, rw_data_type) {
+                                Some(pos)
+                            } else {
+                                None
+                            }
+                        })
+                })
+                .collect();
+            Ok(valid_column_indices)
+        }
+        None => Ok(vec![]),
+    }
+}
+
+/// Reads a specified Parquet file and converts its content into a stream of chunks.
+pub async fn read_parquet_file(
+    op: Operator,
+    file_name: String,
+    rw_columns: Option<Vec<Column>>,
+    parser_columns: Option<Vec<SourceColumnDesc>>,
+    batch_size: usize,
+    offset: usize,
+) -> ConnectorResult<
+    Pin<Box<dyn Stream<Item = Result<StreamChunk, crate::error::ConnectorError>> + Send>>,
+> {
+    let mut reader: tokio_util::compat::Compat<opendal::FuturesAsyncReader> = op
+        .reader_with(&file_name)
+        .into_future() // Unlike `rustc`, `try_stream` seems require manual `into_future`.
+        .await?
+        .into_futures_async_read(..)
+        .await?
+        .compat();
+    let parquet_metadata = reader.get_metadata().await.map_err(anyhow::Error::from)?;
+
+    let file_metadata = parquet_metadata.file_metadata();
+    let column_indices = extract_valid_column_indices(rw_columns, file_metadata)?;
+    let projection_mask = ProjectionMask::leaves(file_metadata.schema_descr(), column_indices);
+    // For the Parquet format, we directly convert from a record batch to a stream chunk.
+    // Therefore, the offset of the Parquet file represents the current position in terms of the number of rows read from the file.
+    let record_batch_stream = ParquetRecordBatchStreamBuilder::new(reader)
+        .await?
+        .with_batch_size(batch_size)
+        .with_projection(projection_mask)
+        .with_offset(offset)
+        .build()?;
+    let converted_arrow_schema = parquet_to_arrow_schema(
+        file_metadata.schema_descr(),
+        file_metadata.key_value_metadata(),
+    )
+    .map_err(anyhow::Error::from)?;
+    let columns = match parser_columns {
+        Some(columns) => columns,
+        None => converted_arrow_schema
+            .fields
+            .iter()
+            .enumerate()
+            .map(|(index, field_ref)| {
+                let data_type = IcebergArrowConvert.type_from_field(field_ref).unwrap();
+                SourceColumnDesc::simple(
+                    field_ref.name().clone(),
+                    data_type,
+                    ColumnId::new(index as i32),
+                )
+            })
+            .collect(),
+    };
+
+    let parquet_parser = ParquetParser::new(columns, file_name, offset)?;
+    let msg_stream: Pin<
+        Box<dyn Stream<Item = Result<StreamChunk, crate::error::ConnectorError>> + Send>,
+    > = parquet_parser.into_stream(record_batch_stream);
+    Ok(msg_stream)
+}
+
+pub async fn get_parquet_fields(
+    op: Operator,
+    file_name: String,
+) -> ConnectorResult<risingwave_common::array::arrow::arrow_schema_udf::Fields> {
+    let mut reader: tokio_util::compat::Compat<opendal::FuturesAsyncReader> = op
+        .reader_with(&file_name)
+        .into_future() // Unlike `rustc`, `try_stream` seems require manual `into_future`.
+        .await?
+        .into_futures_async_read(..)
+        .await?
+        .compat();
+    let parquet_metadata = reader.get_metadata().await.map_err(anyhow::Error::from)?;
+
+    let file_metadata = parquet_metadata.file_metadata();
+    let converted_arrow_schema = parquet_to_arrow_schema(
+        file_metadata.schema_descr(),
+        file_metadata.key_value_metadata(),
+    )
+    .map_err(anyhow::Error::from)?;
+    let fields: risingwave_common::array::arrow::arrow_schema_udf::Fields =
+        converted_arrow_schema.fields;
+    Ok(fields)
+}
+
+/// This function checks whether the schema of a Parquet file matches the user defined schema.
+/// It handles the following special cases:
+/// - Arrow's `timestamp(_, None)` types (all four time units) match with RisingWave's `TimeStamp` type.
+/// - Arrow's `timestamp(_, Some)` matches with RisingWave's `TimeStamptz` type.
+/// - Since RisingWave does not have an `UInt` type:
+///   - Arrow's `UInt8` matches with RisingWave's `Int16`.
+///   - Arrow's `UInt16` matches with RisingWave's `Int32`.
+///   - Arrow's `UInt32` matches with RisingWave's `Int64`.
+///   - Arrow's `UInt64` matches with RisingWave's `Decimal`.
+/// - Arrow's `Float16` matches with RisingWave's `Float32`.
+fn is_parquet_schema_match_source_schema(
+    arrow_data_type: &ArrowDateType,
+    rw_data_type: &RwDataType,
+) -> bool {
+    matches!(
+        (arrow_data_type, rw_data_type),
+        (ArrowDateType::Boolean, RwDataType::Boolean)
+            | (
+                ArrowDateType::Int8 | ArrowDateType::Int16 | ArrowDateType::UInt8,
+                RwDataType::Int16
+            )
+            | (
+                ArrowDateType::Int32 | ArrowDateType::UInt16,
+                RwDataType::Int32
+            )
+            | (
+                ArrowDateType::Int64 | ArrowDateType::UInt32,
+                RwDataType::Int64
+            )
+            | (
+                ArrowDateType::UInt64 | ArrowDateType::Decimal128(_, _),
+                RwDataType::Decimal
+            )
+            | (ArrowDateType::Decimal256(_, _), RwDataType::Int256)
+            | (
+                ArrowDateType::Float16 | ArrowDateType::Float32,
+                RwDataType::Float32
+            )
+            | (ArrowDateType::Float64, RwDataType::Float64)
+            | (ArrowDateType::Timestamp(_, None), RwDataType::Timestamp)
+            | (
+                ArrowDateType::Timestamp(_, Some(_)),
+                RwDataType::Timestamptz
+            )
+            | (ArrowDateType::Date32, RwDataType::Date)
+            | (
+                ArrowDateType::Time32(_) | ArrowDateType::Time64(_),
+                RwDataType::Time
+            )
+            | (
+                ArrowDateType::Interval(IntervalUnit::MonthDayNano),
+                RwDataType::Interval
+            )
+            | (
+                ArrowDateType::Utf8 | ArrowDateType::LargeUtf8,
+                RwDataType::Varchar
+            )
+            | (
+                ArrowDateType::Binary | ArrowDateType::LargeBinary,
+                RwDataType::Bytea
+            )
+            | (ArrowDateType::List(_), RwDataType::List(_))
+            | (ArrowDateType::Struct(_), RwDataType::Struct(_))
+            | (ArrowDateType::Map(_, _), RwDataType::Map(_))
+    )
+}
diff --git a/src/connector/src/source/iceberg/parquet_file_reader.rs b/src/connector/src/source/iceberg/parquet_file_reader.rs
deleted file mode 100644
index 5a01f2b0ed84..000000000000
--- a/src/connector/src/source/iceberg/parquet_file_reader.rs
+++ /dev/null
@@ -1,134 +0,0 @@
-// Copyright 2024 RisingWave Labs
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use std::collections::HashMap;
-use std::ops::Range;
-use std::sync::Arc;
-
-use anyhow::anyhow;
-use bytes::Bytes;
-use futures::future::BoxFuture;
-use futures::{FutureExt, TryFutureExt};
-use iceberg::io::{
-    FileIOBuilder, FileMetadata, FileRead, S3_ACCESS_KEY_ID, S3_REGION, S3_SECRET_ACCESS_KEY,
-};
-use iceberg::{Error, ErrorKind};
-use opendal::layers::RetryLayer;
-use opendal::services::S3;
-use opendal::Operator;
-use parquet::arrow::async_reader::AsyncFileReader;
-use parquet::arrow::ParquetRecordBatchStreamBuilder;
-use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
-use url::Url;
-
-pub struct ParquetFileReader<R: FileRead> {
-    meta: FileMetadata,
-    r: R,
-}
-
-impl<R: FileRead> ParquetFileReader<R> {
-    pub fn new(meta: FileMetadata, r: R) -> Self {
-        Self { meta, r }
-    }
-}
-
-impl<R: FileRead> AsyncFileReader for ParquetFileReader<R> {
-    fn get_bytes(&mut self, range: Range<usize>) -> BoxFuture<'_, parquet::errors::Result<Bytes>> {
-        Box::pin(
-            self.r
-                .read(range.start as _..range.end as _)
-                .map_err(|err| parquet::errors::ParquetError::External(Box::new(err))),
-        )
-    }
-
-    fn get_metadata(&mut self) -> BoxFuture<'_, parquet::errors::Result<Arc<ParquetMetaData>>> {
-        async move {
-            let reader = ParquetMetaDataReader::new();
-            let size = self.meta.size as usize;
-            let meta = reader.load_and_finish(self, size).await?;
-
-            Ok(Arc::new(meta))
-        }
-        .boxed()
-    }
-}
-
-pub async fn create_parquet_stream_builder(
-    s3_region: String,
-    s3_access_key: String,
-    s3_secret_key: String,
-    location: String,
-) -> Result<ParquetRecordBatchStreamBuilder<ParquetFileReader<impl FileRead>>, anyhow::Error> {
-    let mut props = HashMap::new();
-    props.insert(S3_REGION, s3_region.clone());
-    props.insert(S3_ACCESS_KEY_ID, s3_access_key.clone());
-    props.insert(S3_SECRET_ACCESS_KEY, s3_secret_key.clone());
-
-    let file_io_builder = FileIOBuilder::new("s3");
-    let file_io = file_io_builder
-        .with_props(props.into_iter())
-        .build()
-        .map_err(|e| anyhow!(e))?;
-    let parquet_file = file_io.new_input(&location).map_err(|e| anyhow!(e))?;
-
-    let parquet_metadata = parquet_file.metadata().await.map_err(|e| anyhow!(e))?;
-    let parquet_reader = parquet_file.reader().await.map_err(|e| anyhow!(e))?;
-    let parquet_file_reader = ParquetFileReader::new(parquet_metadata, parquet_reader);
-
-    ParquetRecordBatchStreamBuilder::new(parquet_file_reader)
-        .await
-        .map_err(|e| anyhow!(e))
-}
-
-pub async fn list_s3_directory(
-    s3_region: String,
-    s3_access_key: String,
-    s3_secret_key: String,
-    dir: String,
-) -> Result<Vec<String>, anyhow::Error> {
-    let url = Url::parse(&dir)?;
-    let bucket = url.host_str().ok_or_else(|| {
-        Error::new(
-            ErrorKind::DataInvalid,
-            format!("Invalid s3 url: {}, missing bucket", dir),
-        )
-    })?;
-
-    let prefix = format!("s3://{}/", bucket);
-    if dir.starts_with(&prefix) {
-        let mut builder = S3::default();
-        builder = builder
-            .region(&s3_region)
-            .access_key_id(&s3_access_key)
-            .secret_access_key(&s3_secret_key)
-            .bucket(bucket);
-        let op = Operator::new(builder)?
-            .layer(RetryLayer::default())
-            .finish();
-
-        op.list(&dir[prefix.len()..])
-            .await
-            .map_err(|e| anyhow!(e))
-            .map(|list| {
-                list.into_iter()
-                    .map(|entry| prefix.to_string() + entry.path())
-                    .collect()
-            })
-    } else {
-        Err(Error::new(
-            ErrorKind::DataInvalid,
-            format!("Invalid s3 url: {}, should start with {}", dir, prefix),
-        ))?
-    }
-}
diff --git a/src/frontend/src/expr/table_function.rs b/src/frontend/src/expr/table_function.rs
index 5f22398cc583..4bd7be784fb0 100644
--- a/src/frontend/src/expr/table_function.rs
+++ b/src/frontend/src/expr/table_function.rs
@@ -17,7 +17,9 @@ use std::sync::{Arc, LazyLock};
 use itertools::Itertools;
 use risingwave_common::array::arrow::IcebergArrowConvert;
 use risingwave_common::types::{DataType, ScalarImpl, StructType};
-use risingwave_connector::source::iceberg::{create_parquet_stream_builder, list_s3_directory};
+use risingwave_connector::source::iceberg::{
+    get_parquet_fields, list_s3_directory, new_s3_operator,
+};
 pub use risingwave_pb::expr::table_function::PbType as TableFunctionType;
 use risingwave_pb::expr::PbTableFunction;
 use thiserror_ext::AsReport;
@@ -181,7 +183,7 @@ impl TableFunction {
 
                 let schema = tokio::task::block_in_place(|| {
                     RUNTIME.block_on(async {
-                        let parquet_stream_builder = create_parquet_stream_builder(
+                        let op = new_s3_operator(
                             eval_args[2].clone(),
                             eval_args[3].clone(),
                             eval_args[4].clone(),
@@ -189,11 +191,18 @@ impl TableFunction {
                                 Some(files) => files[0].clone(),
                                 None => eval_args[5].clone(),
                             },
+                        )?;
+                        let fields = get_parquet_fields(
+                            op,
+                            match files.as_ref() {
+                                Some(files) => files[0].clone(),
+                                None => eval_args[5].clone(),
+                            },
                         )
                         .await?;
 
                         let mut rw_types = vec![];
-                        for field in parquet_stream_builder.schema().fields() {
+                        for field in &fields {
                             rw_types.push((
                                 field.name().to_string(),
                                 IcebergArrowConvert.type_from_field(field)?,


Test patch:
diff --git a/e2e_test/s3/fs_sink.py b/e2e_test/s3/fs_sink.py
index b068b479d8f8..ac99f9b12119 100644
--- a/e2e_test/s3/fs_sink.py
+++ b/e2e_test/s3/fs_sink.py
@@ -10,6 +10,9 @@
 from time import sleep
 from minio import Minio
 from random import uniform
+from time import sleep
+import numpy as np
+import time
 
 def gen_data(file_num, item_num_per_file):
     assert item_num_per_file % 2 == 0, \
@@ -21,6 +24,12 @@ def gen_data(file_num, item_num_per_file):
             'sex': item_id % 2,
             'mark': (-1) ** (item_id % 2),
             'test_int': pa.scalar(1, type=pa.int32()),
+            'test_int8': pa.scalar(1, type=pa.int8()),
+            'test_uint8': pa.scalar(item_id % 256, type=pa.uint8()),  # UInt8
+            'test_uint16': pa.scalar(item_id % 65536, type=pa.uint16()),  # UInt16
+            'test_uint32': pa.scalar(item_id % (2**32), type=pa.uint32()),  # UInt32
+            'test_uint64': pa.scalar(item_id % (2**64), type=pa.uint64()),  # UInt64
+            'test_float_16': pa.scalar(np.float16(4.0), type=pa.float16()),
             'test_real': pa.scalar(4.0, type=pa.float32()),
             'test_double_precision': pa.scalar(5.0, type=pa.float64()),
             'test_varchar': pa.scalar('7', type=pa.string()),
@@ -60,6 +69,12 @@ def _table():
         sex bigint,
         mark bigint,
         test_int int,
+        test_int8 smallint,
+        test_uint8 smallint,
+        test_uint16 int,
+        test_uint32 bigint,
+        test_uint64 decimal,
+        test_float_16 real,
         test_real real,
         test_double_precision double precision,
         test_varchar varchar,
@@ -135,6 +150,12 @@ def _table():
         sex,
         mark,
         test_int,
+        test_int8,
+        test_uint8,
+        test_uint16,
+        test_uint32,
+        test_uint64,
+        test_float_16,
         test_real,
         test_double_precision,
         test_varchar,
@@ -171,6 +192,12 @@ def _table():
         sex bigint,
         mark bigint,
         test_int int,
+        test_int8 smallint,
+        test_uint8 smallint,
+        test_uint16 int,
+        test_uint32 bigint,
+        test_uint64 decimal,
+        test_float_16 real,
         test_real real,
         test_double_precision double precision,
         test_varchar varchar,
@@ -208,6 +235,7 @@ def _table():
     stmt = f'select count(*), sum(id) from test_parquet_sink_table'
     print(f'Execute reading sink files: {stmt}')
 
+    print(f'Create s3 sink json format')
     # Execute a SELECT statement
     cur.execute(f'''CREATE sink test_file_sink_json as select
         id,
@@ -215,6 +243,12 @@ def _table():
         sex,
         mark,
         test_int,
+        test_int8,
+        test_uint8,
+        test_uint16,
+        test_uint32,
+        test_uint64,
+        test_float_16,
         test_real,
         test_double_precision,
         test_varchar,
@@ -251,6 +285,12 @@ def _table():
         sex bigint,
         mark bigint,
         test_int int,
+        test_int8 smallint,
+        test_uint8 smallint,
+        test_uint16 int,
+        test_uint32 bigint,
+        test_uint64 decimal,
+        test_float_16 real,
         test_real real,
         test_double_precision double precision,
         test_varchar varchar,