add support for creating new iceberg tables with partition keys

Author: jayakasadev

e2e_test/iceberg/benches/predicate_pushdown.slt

@@ -19,7 +19,8 @@ CREATE SINK sink1 AS select * from t1 WITH (
     s3.access.key = 'hummockadmin',
     s3.secret.key = 'hummockadmin',
     commit_checkpoint_interval = 1,
-    create_table_if_not_exists = 'true'
+    create_table_if_not_exists = 'true',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_connection.slt

@@ -26,6 +26,7 @@ CREATE SINK sink1 from s1 WITH (
     create_table_if_not_exists = 'true',
     commit_checkpoint_interval = 1,
     primary_key = 'i1,i2',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_predicate_pushdown.slt

@@ -37,7 +37,8 @@ CREATE SINK sink1 AS select * from mv1 WITH (
     s3.access.key = 'hummockadmin',
     s3.secret.key = 'hummockadmin',
     commit_checkpoint_interval = 1,
-    create_table_if_not_exists = 'true'
+    create_table_if_not_exists = 'true',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_select_empty_table.slt

@@ -25,7 +25,8 @@ CREATE SINK sink1 AS select * from mv1 WITH (
     s3.access.key = 'hummockadmin',
     s3.secret.key = 'hummockadmin',
     commit_checkpoint_interval = 1,
-    create_table_if_not_exists = 'true'
+    create_table_if_not_exists = 'true',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_sink_no_partition_append_only_table.slt

@@ -39,7 +39,8 @@ CREATE SINK s6 AS select * from mv6 WITH (
     s3.region = 'us-east-1',
     s3.access.key = 'hummockadmin',
     s3.secret.key = 'hummockadmin',
-    create_table_if_not_exists = 'true'
+    create_table_if_not_exists = 'true',
+    partition_by = 'year(v_date)'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_source_all_delete.slt

@@ -24,6 +24,7 @@ CREATE SINK sink1 AS select * from mv1 WITH (
     create_table_if_not_exists = 'true',
     commit_checkpoint_interval = 3,
     primary_key = 'i1,i2',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_source_equality_delete.slt

@@ -26,6 +26,7 @@ CREATE SINK sink1 AS select * from mv1 WITH (
     s3.secret.key = 'hummockadmin',
     create_table_if_not_exists = 'true',
     primary_key = 'i1,i2',
+    partition_by = 'i1'
 );
 
 statement ok
@@ -122,6 +123,7 @@ CREATE SINK sink2 AS select * from s2 WITH (
     s3.secret.key = 'hummockadmin',
     create_table_if_not_exists = 'true',
     primary_key = 'i2,i1',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_source_explain_for_delete.slt

@@ -26,6 +26,7 @@ CREATE SINK sink1 AS select * from mv1 WITH (
     s3.secret.key = 'hummockadmin',
     create_table_if_not_exists = 'true',
     primary_key = 'i1,i2',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/iceberg/test_case/iceberg_source_position_delete.slt

@@ -24,6 +24,7 @@ CREATE SINK sink1 AS select * from mv1 WITH (
     create_table_if_not_exists = 'true',
     commit_checkpoint_interval = 3,
     primary_key = 'i1,i2',
+    partition_by = 'i1'
 );
 
 statement ok

e2e_test/sink/iceberg_sink.slt

@@ -67,7 +67,8 @@ CREATE SINK s7 AS select mv6.v1 as v1, mv6.v2 as v2, mv6.v3 as v3 from mv6 WITH
     database.name='demo_db',
     table.name='e2e_auto_create_table',
     commit_checkpoint_interval = 1,
-    create_table_if_not_exists = 'true'
+    create_table_if_not_exists = 'true',
+    partition_by = 'v1'
 );
 
 statement ok

src/connector/src/sink/iceberg/mod.rs

@@ -18,12 +18,15 @@ mod prometheus;
 use std::collections::{BTreeMap, HashMap};
 use std::fmt::Debug;
 use std::num::NonZeroU64;
+use std::str::FromStr;
 use std::sync::Arc;
 
 use anyhow::{anyhow, Context};
 use async_trait::async_trait;
 use iceberg::arrow::{arrow_schema_to_schema, schema_to_arrow_schema};
-use iceberg::spec::{DataFile, SerializedDataFile};
+use iceberg::spec::{
+    DataFile, SerializedDataFile, Transform, UnboundPartitionField, UnboundPartitionSpec,
+};
 use iceberg::table::Table;
 use iceberg::transaction::Transaction;
 use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
@@ -47,6 +50,7 @@ use itertools::Itertools;
 use parquet::file::properties::WriterProperties;
 use prometheus::monitored_general_writer::MonitoredGeneralWriterBuilder;
 use prometheus::monitored_position_delete_writer::MonitoredPositionDeleteWriterBuilder;
+use regex::Regex;
 use risingwave_common::array::arrow::arrow_array_iceberg::{Int32Array, RecordBatch};
 use risingwave_common::array::arrow::arrow_schema_iceberg::{
     self, DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef,
@@ -109,6 +113,13 @@ pub struct IcebergConfig {
     #[serde(skip)]
     pub java_catalog_props: HashMap<String, String>,
 
+    #[serde(
+        rename = "partition_by",
+        default,
+        deserialize_with = "deserialize_optional_string_seq_from_string"
+    )]
+    pub partition_by: Option<Vec<String>>,
+
     /// Commit every n(>0) checkpoints, default is 10.
     #[serde(default = "default_commit_checkpoint_interval")]
     #[serde_as(as = "DisplayFromStr")]
@@ -311,9 +322,53 @@ impl IcebergSink {
                 }
             };
 
+            let partition_fields = match &self.config.partition_by {
+                Some(partition_by) => {
+                    let mut partition_fields =
+                        Vec::<UnboundPartitionField>::with_capacity(partition_by.len());
+                    for partition_field in partition_by {
+                        let re = Regex::new(r"\w+(\(\w+\))?").unwrap();
+                        if !re.is_match(partition_field) {
+                            bail!(format!("Invalid partition field: {}", partition_field))
+                        }
+                        let (_, [field1, field2]) = re.captures(partition_field).unwrap().extract();
+                        let mut func = field1.to_owned();
+                        let mut column = field2.to_owned();
+                        if column.is_empty() {
+                            column.replace_range(.., &func);
+                            func.replace_range(.., "identity");
+                        }
+                        let transform = Transform::from_str(&func).unwrap();
+                        if transform == Transform::Unknown {
+                            bail!(format!("Invalid partition field: {}", partition_field))
+                        }
+                        for (pos, col) in self.param.columns.iter().enumerate() {
+                            if col.name == column {
+                                partition_fields.push(
+                                    UnboundPartitionField::builder()
+                                        .source_id(pos as i32)
+                                        .transform(transform)
+                                        .name(column.to_owned())
+                                        .build(),
+                                );
+                                break;
+                            }
+                        }
+                    }
+                    partition_fields
+                }
+                None => Vec::<UnboundPartitionField>::new(),
+            };
+
             let table_creation_builder = TableCreation::builder()
                 .name(self.config.common.table_name.clone())
-                .schema(iceberg_schema);
+                .schema(iceberg_schema)
+                .partition_spec(
+                    UnboundPartitionSpec::builder()
+                        .add_partition_fields(partition_fields)
+                        .unwrap()
+                        .build(),
+                );
 
             let table_creation = match location {
                 Some(location) => table_creation_builder.location(location).build(),
@@ -1402,6 +1457,7 @@ mod test {
             ("connector", "iceberg"),
             ("type", "upsert"),
             ("primary_key", "v1"),
+            ("partition_by", "v1,identity(v2)"),
             ("warehouse.path", "s3://iceberg"),
             ("s3.endpoint", "http://127.0.0.1:9301"),
             ("s3.access.key", "hummockadmin"),
@@ -1445,6 +1501,7 @@ mod test {
             r#type: "upsert".to_owned(),
             force_append_only: false,
             primary_key: Some(vec!["v1".to_owned()]),
+            partition_by: Some(vec!["v1".to_owned(), "identity(v2)".to_owned()]),
             java_catalog_props: [("jdbc.user", "admin"), ("jdbc.password", "123456")]
                 .into_iter()
                 .map(|(k, v)| (k.to_owned(), v.to_owned()))

src/connector/with_options_sink.yaml

@@ -504,6 +504,10 @@ IcebergConfig:
   - name: java_catalog_props
     field_type: HashMap<String,String>
     required: false
+  - name: partition_by
+    field_type: Vec<String>
+    required: false
+    default: Default::default
   - name: commit_checkpoint_interval
     field_type: u64
     comments: Commit every n(>0) checkpoints, default is 10.