feat(connector): introduce parquet file source

ci/scripts/s3-source-test.sh

@@ -32,7 +32,7 @@ echo "--- starting risingwave cluster with connector node"
 risedev ci-start ci-1cn-1fe
 
 echo "--- Run test"
-python3 -m pip install --break-system-packages minio psycopg2-binary opendal
+python3 -m pip install --break-system-packages minio psycopg2-binary opendal pandas
 if [[ -v format_type ]]; then
  python3 e2e_test/s3/"$script" "$format_type"
 else

ci/workflows/main-cron.yml

@@ -478,6 +478,28 @@ steps:
  timeout_in_minutes: 25
  retry: *auto-retry
 
+ - label: "S3_v2 source check on parquet file"
+ key: "s3-v2-source-check-parquet-file"
+ command: "ci/scripts/s3-source-test.sh -p ci-release -s fs_parquet_source.py"
+ if: |
+ !(build.pull_request.labels includes "ci/main-cron/skip-ci") && build.env("CI_STEPS") == null
+ || build.pull_request.labels includes "ci/run-s3-source-tests"
+ || build.env("CI_STEPS") =~ /(^|,)s3-source-tests?(,|$$)/
+ depends_on: build
+ plugins:
+ - seek-oss/aws-sm#v2.3.1:
+ env:
+ S3_SOURCE_TEST_CONF: ci_s3_source_test_aws
+ - docker-compose#v5.1.0:
+ run: rw-build-env
+ config: ci/docker-compose.yml
+ mount-buildkite-agent: true
+ environment:
+ - S3_SOURCE_TEST_CONF
+ - ./ci/plugins/upload-failure-logs
+ timeout_in_minutes: 25
+ retry: *auto-retry
+
  - label: "S3_v2 source batch read on AWS (json parser)"
  key: "s3-v2-source-batch-read-check-aws-json-parser"
  command: "ci/scripts/s3-source-test.sh -p ci-release -s fs_source_batch.py -t json"

e2e_test/s3/fs_parquet_source.py

@@ -0,0 +1,137 @@
+import os
+import sys
+import random
+import psycopg2
+import json
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pandas as pd
+from datetime import datetime, timezone
+from time import sleep
+from minio import Minio
+from random import uniform
+
+def gen_data(file_num, item_num_per_file):
+ assert item_num_per_file % 2 == 0, \
+ f'item_num_per_file should be even to ensure sum(mark) == 0: {item_num_per_file}'
+ return [
+ [{
+ 'id': file_id * item_num_per_file + item_id,
+ 'name': f'{file_id}_{item_id}_{file_id * item_num_per_file + item_id}',
+ 'sex': item_id % 2,
+ 'mark': (-1) ** (item_id % 2),
+ 'test_int': pa.scalar(1, type=pa.int32()),
+ '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()),
+ 'test_bytea': pa.scalar(b'\xDe00BeEf', type=pa.binary()),
+ 'test_date': pa.scalar(datetime.now().date(), type=pa.date32()),
+ 'test_time': pa.scalar(datetime.now().time(), type=pa.time64('us')),
+ 'test_timestamp': pa.scalar(datetime.now().timestamp() * 1000000, type=pa.timestamp('us')),
+ 'test_timestamptz': pa.scalar(datetime.now().timestamp() * 1000, type=pa.timestamp('us', tz='+00:00')),
+ } for item_id in range(item_num_per_file)]
+ for file_id in range(file_num)
+ ]
+
+def do_test(config, file_num, item_num_per_file, prefix):
+ conn = psycopg2.connect(
+ host="localhost",
+ port="4566",
+ user="root",
+ database="dev"
+ )
+
+ # Open a cursor to execute SQL statements
+ cur = conn.cursor()
+
+ def _table():
+ return 's3_test_parquet'
+
+ # Execute a SELECT statement
+ cur.execute(f'''CREATE TABLE {_table()}(
+ id bigint primary key,
+ name TEXT,
+ sex bigint,
+ mark bigint,
+ test_int int,
+ test_real real,
+ test_double_precision double precision,
+ test_varchar varchar,
+ test_bytea bytea,
+ test_date date,
+ test_time time,
+ test_timestamp timestamp,
+ test_timestamptz timestamptz,
+ ) WITH (
+ connector = 's3_v2',
+ match_pattern = '*.parquet',
+ s3.region_name = '{config['S3_REGION']}',
+ s3.bucket_name = '{config['S3_BUCKET']}',
+ s3.credentials.access = '{config['S3_ACCESS_KEY']}',
+ s3.credentials.secret = '{config['S3_SECRET_KEY']}',
+ s3.endpoint_url = 'https://{config['S3_ENDPOINT']}'
+ ) FORMAT PLAIN ENCODE PARQUET;''')
+
+ total_rows = file_num * item_num_per_file
+ MAX_RETRIES = 40
+ for retry_no in range(MAX_RETRIES):
+ cur.execute(f'select count(*) from {_table()}')
+ result = cur.fetchone()
+ if result[0] == total_rows:
+ break
+ print(f"[retry {retry_no}] Now got {result[0]} rows in table, {total_rows} expected, wait 10s")
+ sleep(10)
+
+ stmt = f'select count(*), sum(id) from {_table()}'
+ print(f'Execute {stmt}')
+ cur.execute(stmt)
+ result = cur.fetchone()
+
+ print('Got:', result)
+
+ def _assert_eq(field, got, expect):
+ assert got == expect, f'{field} assertion failed: got {got}, expect {expect}.'
+
+ _assert_eq('count(*)', result[0], total_rows)
+ _assert_eq('sum(id)', result[1], (total_rows - 1) * total_rows / 2)
+
+ print('Test pass')
+
+ cur.execute(f'drop table {_table()}')
+ cur.close()
+ conn.close()
+
+
+if __name__ == "__main__":
+ FILE_NUM = 10
+ ITEM_NUM_PER_FILE = 2000
+ data = gen_data(FILE_NUM, ITEM_NUM_PER_FILE)
+
+ config = json.loads(os.environ["S3_SOURCE_TEST_CONF"])
+ client = Minio(
+ config["S3_ENDPOINT"],
+ access_key=config["S3_ACCESS_KEY"],
+ secret_key=config["S3_SECRET_KEY"],
+ secure=True,
+ )
+ run_id = str(random.randint(1000, 9999))
+ _local = lambda idx: f'data_{idx}.parquet'
+ _s3 = lambda idx: f"{run_id}_data_{idx}.parquet"
+
+ # put s3 files
+ for idx, file_data in enumerate(data):
+ table = pa.Table.from_pandas(pd.DataFrame(file_data))
+ pq.write_table(table, _local(idx))
+
+ client.fput_object(
+ config["S3_BUCKET"],
+ _s3(idx),
+ _local(idx)
+ )
+
+ # do test
+ do_test(config, FILE_NUM, ITEM_NUM_PER_FILE, run_id)
+
+ # clean up s3 files
+ for idx, _ in enumerate(data):
+ client.remove_object(config["S3_BUCKET"], _s3(idx))

proto/plan_common.proto

@@ -150,6 +150,7 @@ enum EncodeType {
  ENCODE_TYPE_TEMPLATE = 7;
  ENCODE_TYPE_NONE = 8;
  ENCODE_TYPE_TEXT = 9;
+ ENCODE_TYPE_PARQUET = 10;
 }
 
 enum RowFormatType {

src/common/src/array/arrow/arrow_impl.rs

@@ -524,9 +524,12 @@ pub trait FromArrow {
  Float64 => self.from_float64_array(array.as_any().downcast_ref().unwrap()),
  Date32 => self.from_date32_array(array.as_any().downcast_ref().unwrap()),
  Time64(Microsecond) => self.from_time64us_array(array.as_any().downcast_ref().unwrap()),
- Timestamp(Microsecond, _) => {
+ Timestamp(Microsecond, None) => {
  self.from_timestampus_array(array.as_any().downcast_ref().unwrap())
  }
+ Timestamp(Microsecond, Some(_)) => {
+ self.from_timestampus_some_array(array.as_any().downcast_ref().unwrap())
+ }
  Interval(MonthDayNano) => {
  self.from_interval_array(array.as_any().downcast_ref().unwrap())
  }
@@ -628,6 +631,13 @@ pub trait FromArrow {
  Ok(ArrayImpl::Timestamp(array.into()))
  }
 
+ fn from_timestampus_some_array(
+ &self,
+ array: &arrow_array::TimestampMicrosecondArray,
+ ) -> Result<ArrayImpl, ArrayError> {
+ Ok(ArrayImpl::Timestamptz(array.into()))
+ }
+
  fn from_interval_array(
  &self,
  array: &arrow_array::IntervalMonthDayNanoArray,

src/connector/Cargo.toml

@@ -93,6 +93,7 @@ opendal = { workspace = true, features = [
 ] }
 openssl = "0.10"
 parking_lot = { workspace = true }
+parquet = { workspace = true }
 paste = "1"
 pg_bigdecimal = { git = "https:/risingwavelabs/rust-pg_bigdecimal", rev = "0b7893d88894ca082b4525f94f812da034486f7c" }
 postgres-openssl = "0.5.0"

src/connector/src/error.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use risingwave_common::array::ArrayError;
 use risingwave_common::error::def_anyhow_newtype;
 use risingwave_pb::PbFieldNotFound;
 use risingwave_rpc_client::error::RpcError;
@@ -41,17 +42,20 @@ def_anyhow_newtype! {
  url::ParseError => "failed to parse url",
  serde_json::Error => "failed to parse json",
  csv::Error => "failed to parse csv",
+
  uuid::Error => transparent, // believed to be self-explanatory
 
  // Connector errors
  opendal::Error => transparent, // believed to be self-explanatory
-
+ parquet::errors::ParquetError => transparent,
+ ArrayError => "Array error",
  sqlx::Error => transparent, // believed to be self-explanatory
  mysql_async::Error => "MySQL error",
  tokio_postgres::Error => "Postgres error",
  apache_avro::Error => "Avro error",
  rdkafka::error::KafkaError => "Kafka error",
  pulsar::Error => "Pulsar error",
+
  async_nats::jetstream::consumer::StreamError => "Nats error",
  async_nats::jetstream::consumer::pull::MessagesError => "Nats error",
  async_nats::jetstream::context::CreateStreamError => "Nats error",

src/connector/src/parser/mod.rs

@@ -24,6 +24,7 @@ pub use debezium::*;
 use futures::{Future, TryFutureExt};
 use futures_async_stream::try_stream;
 pub use json_parser::*;
+pub use parquet_parser::ParquetParser;
 pub use protobuf::*;
 use risingwave_common::array::{ArrayBuilderImpl, Op, StreamChunk};
 use risingwave_common::bail;
@@ -76,6 +77,7 @@ mod debezium;
 mod json_parser;
 mod maxwell;
 mod mysql;
+pub mod parquet_parser;
 pub mod plain_parser;
 mod postgres;
 
@@ -1117,6 +1119,7 @@ pub enum EncodingProperties {
  Json(JsonProperties),
  MongoJson,
  Bytes(BytesProperties),
+ Parquet,
  Native,
  /// Encoding can't be specified because the source will determines it. Now only used in Iceberg.
  None,
@@ -1170,6 +1173,7 @@ impl SpecificParserConfig {
  delimiter: info.csv_delimiter as u8,
  has_header: info.csv_has_header,
  }),
+ (SourceFormat::Plain, SourceEncode::Parquet) => EncodingProperties::Parquet,
  (SourceFormat::Plain, SourceEncode::Avro)
  | (SourceFormat::Upsert, SourceEncode::Avro) => {
  let mut config = AvroProperties {