risingwave_connector/sink/iceberg/
commit.rs

1// Copyright 2026 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;
16use std::time::Duration;
17
18use anyhow::{Context, anyhow};
19use async_trait::async_trait;
20use iceberg::Catalog;
21use iceberg::arrow::schema_to_arrow_schema;
22use iceberg::spec::{DataFile, Operation, SerializedDataFile, TableMetadata};
23use iceberg::table::Table;
24use iceberg::transaction::{ApplyTransactionAction, FastAppendAction, Transaction};
25use itertools::Itertools;
26use risingwave_common::array::arrow::arrow_schema_iceberg::{
27    DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields,
28};
29use risingwave_common::array::arrow::{IcebergArrowConvert, IcebergCreateTableArrowConvert};
30use risingwave_common::bail;
31use risingwave_common::catalog::Field;
32use risingwave_common::error::IcebergError;
33use risingwave_pb::connector_service::SinkMetadata;
34use risingwave_pb::connector_service::sink_metadata::Metadata::Serialized;
35use risingwave_pb::connector_service::sink_metadata::SerializedMetadata;
36use risingwave_pb::stream_plan::PbSinkSchemaChange;
37use serde::{Deserialize, Serialize};
38use serde_json::from_value;
39use thiserror_ext::AsReport;
40use tokio::sync::mpsc::UnboundedSender;
41use tracing::warn;
42
43use super::commit_retry::{self, CommitError};
44use super::{GLOBAL_SINK_METRICS, IcebergConfig, SinkError, commit_branch, resolve_partition_type};
45use crate::connector_common::{IcebergCommittedSnapshot, IcebergSinkCompactionUpdate};
46use crate::sink::catalog::SinkId;
47use crate::sink::{Result, SinglePhaseCommitCoordinator, SinkParam, TwoPhaseCommitCoordinator};
48
49const SCHEMA_ID: &str = "schema_id";
50const PARTITION_SPEC_ID: &str = "partition_spec_id";
51const DATA_FILES: &str = "data_files";
52
53#[derive(Default, Clone)]
54pub struct IcebergCommitResult {
55    pub schema_id: i32,
56    pub partition_spec_id: i32,
57    pub data_files: Vec<SerializedDataFile>,
58}
59
60impl IcebergCommitResult {
61    pub fn try_from(value: &SinkMetadata) -> Result<Self> {
62        let Some(Serialized(value)) = &value.metadata else {
63            bail!("Can't create iceberg sink write result from empty data!");
64        };
65
66        Self::try_from_serialized_bytes(&value.metadata)
67    }
68
69    pub fn try_from_serialized_bytes(value: &[u8]) -> Result<Self> {
70        let mut values = if let serde_json::Value::Object(value) =
71            serde_json::from_slice::<serde_json::Value>(value)
72                .context("Can't parse iceberg sink metadata")?
73        {
74            value
75        } else {
76            bail!("iceberg sink metadata should be an object");
77        };
78
79        let schema_id;
80        if let Some(serde_json::Value::Number(value)) = values.remove(SCHEMA_ID) {
81            schema_id = value
82                .as_u64()
83                .ok_or_else(|| anyhow!("schema_id should be a u64"))?;
84        } else {
85            bail!("iceberg sink metadata should have schema_id");
86        }
87
88        let partition_spec_id;
89        if let Some(serde_json::Value::Number(value)) = values.remove(PARTITION_SPEC_ID) {
90            partition_spec_id = value
91                .as_u64()
92                .ok_or_else(|| anyhow!("partition_spec_id should be a u64"))?;
93        } else {
94            bail!("iceberg sink metadata should have partition_spec_id");
95        }
96
97        let data_files: Vec<SerializedDataFile>;
98        if let serde_json::Value::Array(values) = values
99            .remove(DATA_FILES)
100            .ok_or_else(|| anyhow!("iceberg sink metadata should have data_files object"))?
101        {
102            data_files = values
103                .into_iter()
104                .map(from_value::<SerializedDataFile>)
105                .collect::<std::result::Result<_, _>>()
106                .unwrap();
107        } else {
108            bail!("iceberg sink metadata should have data_files object");
109        }
110
111        Ok(Self {
112            schema_id: schema_id as i32,
113            partition_spec_id: partition_spec_id as i32,
114            data_files,
115        })
116    }
117}
118
119impl<'a> TryFrom<&'a IcebergCommitResult> for SinkMetadata {
120    type Error = SinkError;
121
122    fn try_from(value: &'a IcebergCommitResult) -> std::result::Result<SinkMetadata, Self::Error> {
123        let bytes = <Vec<u8>>::try_from(value)?;
124        Ok(SinkMetadata {
125            metadata: Some(Serialized(SerializedMetadata { metadata: bytes })),
126        })
127    }
128}
129
130impl<'a> TryFrom<&'a IcebergCommitResult> for Vec<u8> {
131    type Error = SinkError;
132
133    fn try_from(value: &'a IcebergCommitResult) -> std::result::Result<Vec<u8>, Self::Error> {
134        let json_data_files = serde_json::Value::Array(
135            value
136                .data_files
137                .iter()
138                .map(serde_json::to_value)
139                .collect::<std::result::Result<Vec<serde_json::Value>, _>>()
140                .context("Can't serialize data files to json")?,
141        );
142        let json_value = serde_json::Value::Object(
143            vec![
144                (
145                    SCHEMA_ID.to_owned(),
146                    serde_json::Value::Number(value.schema_id.into()),
147                ),
148                (
149                    PARTITION_SPEC_ID.to_owned(),
150                    serde_json::Value::Number(value.partition_spec_id.into()),
151                ),
152                (DATA_FILES.to_owned(), json_data_files),
153            ]
154            .into_iter()
155            .collect(),
156        );
157        Ok(serde_json::to_vec(&json_value).context("Can't serialize iceberg sink metadata")?)
158    }
159}
160
161#[derive(Default, Clone, Serialize, Deserialize)]
162pub struct IcebergPositionDeleteMergerCommitResult {
163    pub schema_id: i32,
164    pub partition_spec_id: i32,
165    pub delete_files: Vec<SerializedDataFile>,
166    pub overwrite_files: Vec<SerializedDataFile>,
167}
168
169impl<'a> TryFrom<&'a SinkMetadata> for IcebergPositionDeleteMergerCommitResult {
170    type Error = SinkError;
171
172    fn try_from(value: &'a SinkMetadata) -> Result<Self> {
173        let Some(Serialized(value)) = &value.metadata else {
174            bail!("Can't create iceberg dv merger commit result from empty data!");
175        };
176        let value = serde_json::from_slice(&value.metadata)
177            .context("Can't deserialize iceberg dv merger commit result from metadata")?;
178        Ok(value)
179    }
180}
181
182impl<'a> TryFrom<&'a IcebergPositionDeleteMergerCommitResult> for SinkMetadata {
183    type Error = SinkError;
184
185    fn try_from(value: &'a IcebergPositionDeleteMergerCommitResult) -> Result<SinkMetadata> {
186        let bytes = serde_json::to_vec(value)
187            .context("Can't serialize iceberg dv merger commit result to metadata")?;
188        Ok(SinkMetadata {
189            metadata: Some(Serialized(SerializedMetadata { metadata: bytes })),
190        })
191    }
192}
193
194fn arrow_data_type_compatible(current: &ArrowDataType, expected: &ArrowDataType) -> bool {
195    use ArrowDataType::*;
196
197    match (current, expected) {
198        // RW Decimal has no precision/scale. Sink creation already accepts any
199        // table Decimal128 precision/scale, while expected schemas here are
200        // generated from RW columns using the same canonical mapping. Ignoring
201        // both values prevents false schema-state ambiguity and cannot mask an
202        // auto schema change, which only supports add/drop columns.
203        (Decimal128(_, _), Decimal128(_, _)) => true,
204        (Binary, LargeBinary) | (LargeBinary, Binary) => true,
205        (List(current_field), List(expected_field)) => {
206            arrow_data_type_compatible(current_field.data_type(), expected_field.data_type())
207        }
208        (Map(current_field, current_sorted), Map(expected_field, expected_sorted)) => {
209            current_sorted == expected_sorted
210                && arrow_data_type_compatible(current_field.data_type(), expected_field.data_type())
211        }
212        (Struct(current_fields), Struct(expected_fields)) => {
213            let expected_fields = expected_fields
214                .iter()
215                .map(|field| field.as_ref().clone())
216                .collect_vec();
217            schema_contains_same_fields(current_fields, &expected_fields)
218        }
219        _ => current == expected,
220    }
221}
222
223fn schema_contains_same_fields(current: &ArrowFields, expected: &[ArrowField]) -> bool {
224    if current.len() != expected.len() {
225        return false;
226    }
227
228    let mut unmatched_current = current.iter().collect_vec();
229    expected.iter().all(|expected_field| {
230        let Some(pos) = unmatched_current.iter().position(|current_field| {
231            current_field.name() == expected_field.name()
232                && arrow_data_type_compatible(current_field.data_type(), expected_field.data_type())
233        }) else {
234            return false;
235        };
236        unmatched_current.swap_remove(pos);
237        true
238    })
239}
240
241pub struct IcebergSinkCommitter {
242    pub(super) catalog: Arc<dyn Catalog>,
243    pub(super) table: Table,
244    pub last_commit_epoch: u64,
245    pub(crate) sink_id: SinkId,
246    pub(crate) config: IcebergConfig,
247    pub(crate) param: SinkParam,
248    pub(super) commit_retry_num: u32,
249    pub(crate) iceberg_compact_stat_sender: Option<UnboundedSender<IcebergSinkCompactionUpdate>>,
250}
251
252impl IcebergSinkCommitter {
253    fn latest_observed_snapshot(&self) -> Option<IcebergCommittedSnapshot> {
254        let branch = commit_branch(self.config.r#type.as_str(), self.config.write_mode);
255        self.table
256            .metadata()
257            .snapshot_for_ref(&branch)
258            .map(|snapshot| IcebergCommittedSnapshot {
259                branch,
260                snapshot_id: snapshot.snapshot_id(),
261                timestamp_ms: snapshot.timestamp_ms(),
262            })
263    }
264
265    fn notify_iceberg_compaction_scheduler(&self, force_compaction: bool) {
266        let Some(iceberg_compact_stat_sender) = &self.iceberg_compact_stat_sender else {
267            return;
268        };
269
270        let Some(observed_snapshot) = self.latest_observed_snapshot() else {
271            warn!(
272                sink_id = %self.sink_id,
273                "skip iceberg compaction update because no observed snapshot is available"
274            );
275            return;
276        };
277
278        let observed_snapshot_id = observed_snapshot.snapshot_id;
279        let observed_snapshot_timestamp_ms = observed_snapshot.timestamp_ms;
280        let observed_snapshot_branch = observed_snapshot.branch.clone();
281
282        if iceberg_compact_stat_sender
283            .send(IcebergSinkCompactionUpdate {
284                sink_id: self.sink_id,
285                force_compaction,
286                observed_snapshot,
287            })
288            .is_err()
289        {
290            warn!(
291                sink_id = %self.sink_id,
292                force_compaction,
293                observed_snapshot_id,
294                observed_snapshot_timestamp_ms,
295                observed_snapshot_branch = %observed_snapshot_branch,
296                "failed to send iceberg compaction update"
297            );
298        }
299    }
300}
301
302#[async_trait]
303impl SinglePhaseCommitCoordinator for IcebergSinkCommitter {
304    async fn init(&mut self) -> Result<()> {
305        tracing::info!(
306            sink_id = %self.param.sink_id,
307            "Iceberg sink coordinator initialized",
308        );
309
310        Ok(())
311    }
312
313    async fn commit_data(&mut self, epoch: u64, metadata: Vec<SinkMetadata>) -> Result<()> {
314        tracing::debug!("Starting iceberg direct commit in epoch {epoch}");
315
316        if metadata.is_empty() {
317            tracing::debug!(?epoch, "No datafile to commit");
318            return Ok(());
319        }
320
321        // Commit data if present
322        if let Some((write_results, snapshot_id)) = self.pre_commit_inner(epoch, metadata)? {
323            self.commit_data_impl(epoch, write_results, snapshot_id)
324                .await?;
325        }
326
327        Ok(())
328    }
329
330    async fn commit_schema_change(
331        &mut self,
332        epoch: u64,
333        schema_change: PbSinkSchemaChange,
334    ) -> Result<()> {
335        tracing::info!(
336            "Committing schema change {:?} in epoch {}",
337            schema_change,
338            epoch
339        );
340        self.commit_schema_change_impl(schema_change).await?;
341        tracing::info!("Successfully committed schema change in epoch {}", epoch);
342
343        Ok(())
344    }
345}
346
347#[async_trait]
348impl TwoPhaseCommitCoordinator for IcebergSinkCommitter {
349    async fn init(&mut self) -> Result<()> {
350        tracing::info!(
351            sink_id = %self.param.sink_id,
352            "Iceberg sink coordinator initialized",
353        );
354
355        Ok(())
356    }
357
358    async fn pre_commit(
359        &mut self,
360        epoch: u64,
361        metadata: Vec<SinkMetadata>,
362        _schema_change: Option<PbSinkSchemaChange>,
363    ) -> Result<Option<Vec<u8>>> {
364        tracing::debug!("Starting iceberg pre commit in epoch {epoch}");
365
366        let (write_results, snapshot_id) = match self.pre_commit_inner(epoch, metadata)? {
367            Some((write_results, snapshot_id)) => (write_results, snapshot_id),
368            None => {
369                tracing::debug!(?epoch, "no data to pre commit");
370                return Ok(None);
371            }
372        };
373
374        let mut write_results_bytes = Vec::new();
375        for each_parallelism_write_result in write_results {
376            let each_parallelism_write_result_bytes =
377                <Vec<u8>>::try_from(&each_parallelism_write_result)?;
378            write_results_bytes.push(each_parallelism_write_result_bytes);
379        }
380
381        let snapshot_id_bytes: Vec<u8> = snapshot_id.to_le_bytes().to_vec();
382        write_results_bytes.push(snapshot_id_bytes);
383
384        let pre_commit_metadata_bytes: Vec<u8> = serialize_metadata(write_results_bytes);
385        Ok(Some(pre_commit_metadata_bytes))
386    }
387
388    async fn commit_data(&mut self, epoch: u64, commit_metadata: Vec<u8>) -> Result<()> {
389        tracing::debug!("Starting iceberg commit in epoch {epoch}");
390
391        if commit_metadata.is_empty() {
392            tracing::debug!(?epoch, "No datafile to commit");
393            return Ok(());
394        }
395
396        // Deserialize commit metadata
397        let mut payload = deserialize_metadata(commit_metadata);
398        if payload.is_empty() {
399            return Err(SinkError::Iceberg(anyhow!(
400                "Invalid commit metadata: empty payload"
401            )));
402        }
403
404        // Last element is snapshot_id
405        let snapshot_id_bytes = payload.pop().ok_or_else(|| {
406            SinkError::Iceberg(anyhow!("Invalid commit metadata: missing snapshot_id"))
407        })?;
408        let snapshot_id = i64::from_le_bytes(
409            snapshot_id_bytes
410                .try_into()
411                .map_err(|_| SinkError::Iceberg(anyhow!("Invalid snapshot id bytes")))?,
412        );
413
414        // Remaining elements are write_results
415        let write_results = payload
416            .into_iter()
417            .map(|p| IcebergCommitResult::try_from_serialized_bytes(&p))
418            .collect::<Result<Vec<_>>>()?;
419
420        let snapshot_committed = self
421            .is_snapshot_id_in_iceberg(&self.config, snapshot_id)
422            .await?;
423
424        if snapshot_committed {
425            tracing::info!(
426                "Snapshot id {} already committed in iceberg table, skip committing again.",
427                snapshot_id
428            );
429            return Ok(());
430        }
431
432        self.commit_data_impl(epoch, write_results, snapshot_id)
433            .await
434    }
435
436    async fn commit_schema_change(
437        &mut self,
438        epoch: u64,
439        schema_change: PbSinkSchemaChange,
440    ) -> Result<()> {
441        let schema_updated = self.check_schema_change_applied(&schema_change)?;
442        if schema_updated {
443            tracing::info!("Schema change already committed in epoch {}, skip", epoch);
444            return Ok(());
445        }
446
447        tracing::info!(
448            "Committing schema change {:?} in epoch {}",
449            schema_change,
450            epoch
451        );
452        self.commit_schema_change_impl(schema_change).await?;
453        tracing::info!("Successfully committed schema change in epoch {epoch}");
454
455        Ok(())
456    }
457
458    async fn abort(&mut self, _epoch: u64, _commit_metadata: Vec<u8>) {
459        // TODO: Files that have been written but not committed should be deleted.
460        tracing::debug!("Abort not implemented yet");
461    }
462}
463
464/// Methods Required to Achieve Exactly Once Semantics
465impl IcebergSinkCommitter {
466    fn pre_commit_inner(
467        &mut self,
468        _epoch: u64,
469        metadata: Vec<SinkMetadata>,
470    ) -> Result<Option<(Vec<IcebergCommitResult>, i64)>> {
471        let write_results: Vec<IcebergCommitResult> = metadata
472            .iter()
473            .map(IcebergCommitResult::try_from)
474            .collect::<Result<Vec<IcebergCommitResult>>>()?;
475
476        // Skip if no data to commit
477        if write_results.is_empty() || write_results.iter().all(|r| r.data_files.is_empty()) {
478            return Ok(None);
479        }
480
481        let expect_schema_id = write_results[0].schema_id;
482        let expect_partition_spec_id = write_results[0].partition_spec_id;
483
484        // guarantee that all write results has same schema_id and partition_spec_id
485        if write_results
486            .iter()
487            .any(|r| r.schema_id != expect_schema_id)
488            || write_results
489                .iter()
490                .any(|r| r.partition_spec_id != expect_partition_spec_id)
491        {
492            return Err(SinkError::Iceberg(anyhow!(
493                "schema_id and partition_spec_id should be the same in all write results"
494            )));
495        }
496
497        let snapshot_id = FastAppendAction::generate_snapshot_id(&self.table);
498
499        Ok(Some((write_results, snapshot_id)))
500    }
501
502    async fn commit_data_impl(
503        &mut self,
504        epoch: u64,
505        write_results: Vec<IcebergCommitResult>,
506        snapshot_id: i64,
507    ) -> Result<()> {
508        // Empty write results should be handled before calling this function.
509        assert!(
510            !write_results.is_empty() && !write_results.iter().all(|r| r.data_files.is_empty())
511        );
512
513        // Check snapshot limit before proceeding with commit
514        self.wait_for_snapshot_limit().await?;
515
516        let expect_schema_id = write_results[0].schema_id;
517        let expect_partition_spec_id = write_results[0].partition_spec_id;
518
519        // Load the latest table to avoid concurrent modification with the best effort.
520        self.table = commit_retry::reload_table(
521            self.catalog.as_ref(),
522            self.table.identifier(),
523            expect_schema_id,
524            expect_partition_spec_id,
525        )
526        .await
527        .map_err(SinkError::Iceberg)?;
528
529        let Some(schema) = self.table.metadata().schema_by_id(expect_schema_id) else {
530            return Err(SinkError::Iceberg(anyhow!(
531                "Can't find schema by id {}",
532                expect_schema_id
533            )));
534        };
535        let partition_type = resolve_partition_type(&self.table, expect_partition_spec_id, schema)?;
536
537        let data_files = write_results
538            .into_iter()
539            .flat_map(|r| {
540                r.data_files.into_iter().map(|f| {
541                    f.try_into(expect_partition_spec_id, &partition_type, schema)
542                        .map_err(|err| SinkError::Iceberg(anyhow!(err)))
543                })
544            })
545            .collect::<Result<Vec<DataFile>>>()?;
546
547        // # TODO:
548        // This retry behavior should be revert and do in iceberg-rust when it supports retry(Track in: https://github.com/apache/iceberg-rust/issues/964)
549        // because retry logic involved reapply the commit metadata.
550        // For now, we just retry the commit operation.
551        let catalog = self.catalog.clone();
552        let table_ident = self.table.identifier().clone();
553        let target_branch = commit_branch(self.config.r#type.as_str(), self.config.write_mode);
554
555        let table = commit_retry::run_with_retry(
556            catalog.clone(),
557            table_ident.clone(),
558            expect_schema_id,
559            expect_partition_spec_id,
560            self.commit_retry_num as usize,
561            |table| {
562                let target_branch = target_branch.clone();
563                let data_files = data_files.clone();
564                let catalog = catalog.clone();
565                async move {
566                    let txn = Transaction::new(&table);
567                    let append_action = txn
568                        .fast_append()
569                        .set_snapshot_id(snapshot_id)
570                        .set_target_branch(target_branch)
571                        .add_data_files(data_files);
572
573                    let tx = append_action.apply(txn).map_err(|err| {
574                        let err: IcebergError = err.into();
575                        tracing::error!(error = %err.as_report(), "Failed to apply iceberg fast_append action");
576                        CommitError::Commit(anyhow!(err).context("apply iceberg fast_append"))
577                    })?;
578
579                    let table = tx.commit(catalog.as_ref()).await.map_err(|err| {
580                        let err: IcebergError = err.into();
581                        tracing::error!(error = %err.as_report(), "Failed to commit iceberg table");
582                        CommitError::Commit(anyhow!(err).context("commit iceberg transaction"))
583                    })?;
584                    Ok(table)
585                }
586            },
587        )
588        .await
589        .map_err(SinkError::Iceberg)?;
590        self.table = table;
591
592        let snapshot_num = self.table.metadata().snapshots().count();
593        let catalog_name = self.config.common.catalog_name();
594        let table_name = self.table.identifier().to_string();
595        let metrics_labels = [&self.param.sink_name, &catalog_name, &table_name];
596        GLOBAL_SINK_METRICS
597            .iceberg_snapshot_num
598            .with_guarded_label_values(&metrics_labels)
599            .set(snapshot_num as i64);
600
601        tracing::debug!("Succeeded to commit to iceberg table in epoch {epoch}.");
602
603        self.notify_iceberg_compaction_scheduler(false);
604
605        Ok(())
606    }
607
608    /// During pre-commit metadata, we record the `snapshot_id` corresponding to each batch of files.
609    /// Therefore, the logic for checking whether all files in this batch are present in Iceberg
610    /// has been changed to verifying if their corresponding `snapshot_id` exists in Iceberg.
611    async fn is_snapshot_id_in_iceberg(
612        &self,
613        iceberg_config: &IcebergConfig,
614        snapshot_id: i64,
615    ) -> Result<bool> {
616        let table = iceberg_config.load_table().await?;
617        if table.metadata().snapshot_by_id(snapshot_id).is_some() {
618            Ok(true)
619        } else {
620            Ok(false)
621        }
622    }
623
624    /// Check if the specified columns already exist in the iceberg table's current schema.
625    /// This is used to determine if schema change has already been applied.
626    fn check_schema_change_applied(&self, schema_change: &PbSinkSchemaChange) -> Result<bool> {
627        let current_schema = self.table.metadata().current_schema();
628        let current_arrow_schema = schema_to_arrow_schema(current_schema.as_ref())
629            .context("Failed to convert schema")
630            .map_err(SinkError::Iceberg)?;
631
632        let iceberg_arrow_convert = IcebergArrowConvert;
633
634        let schema_matches = |expected: &[ArrowField]| {
635            schema_contains_same_fields(current_arrow_schema.fields(), expected)
636        };
637
638        let original_arrow_fields: Vec<ArrowField> = schema_change
639            .original_schema
640            .iter()
641            .map(|pb_field| {
642                let field = Field::from(pb_field);
643                iceberg_arrow_convert
644                    .to_arrow_field(&field.name, &field.data_type)
645                    .context("Failed to convert field to arrow")
646                    .map_err(SinkError::Iceberg)
647            })
648            .collect::<Result<_>>()?;
649
650        // If current schema equals original_schema, then schema change is NOT applied.
651        if schema_matches(&original_arrow_fields) {
652            tracing::debug!(
653                "Current iceberg schema matches original_schema ({} columns); schema change not applied",
654                original_arrow_fields.len()
655            );
656            return Ok(false);
657        }
658
659        let expected_after_change = match schema_change.op.as_ref() {
660            Some(risingwave_pb::stream_plan::sink_schema_change::Op::AddColumns(
661                add_columns_op,
662            )) => {
663                let add_arrow_fields: Vec<ArrowField> = add_columns_op
664                    .fields
665                    .iter()
666                    .map(|pb_field| {
667                        let field = Field::from(pb_field);
668                        iceberg_arrow_convert
669                            .to_arrow_field(&field.name, &field.data_type)
670                            .context("Failed to convert field to arrow")
671                            .map_err(SinkError::Iceberg)
672                    })
673                    .collect::<Result<_>>()?;
674
675                let mut expected_after_change = original_arrow_fields;
676                expected_after_change.extend(add_arrow_fields);
677                expected_after_change
678            }
679            Some(risingwave_pb::stream_plan::sink_schema_change::Op::DropColumns(
680                drop_columns_op,
681            )) => original_arrow_fields
682                .into_iter()
683                .filter(|field| {
684                    !drop_columns_op
685                        .column_names
686                        .iter()
687                        .any(|name| name == field.name())
688                })
689                .collect_vec(),
690            _ => {
691                return Err(SinkError::Iceberg(anyhow!(
692                    "Unsupported sink schema change op in iceberg sink: {:?}",
693                    schema_change.op
694                )));
695            }
696        };
697
698        // If current schema equals the changed schema, then schema change is applied.
699        if schema_matches(&expected_after_change) {
700            tracing::debug!(
701                "Current iceberg schema matches changed schema ({} columns); schema change already applied",
702                expected_after_change.len()
703            );
704            return Ok(true);
705        }
706
707        Err(SinkError::Iceberg(anyhow!(
708            "Current iceberg schema does not match either original_schema ({} cols) or changed schema; cannot determine whether schema change is applied",
709            schema_change.original_schema.len()
710        )))
711    }
712
713    /// Commit schema changes (e.g., add columns) to the iceberg table.
714    /// This function uses Transaction API to atomically update the table schema
715    /// with optimistic locking to prevent concurrent conflicts.
716    async fn commit_schema_change_impl(&mut self, schema_change: PbSinkSchemaChange) -> Result<()> {
717        use iceberg::spec::NestedField;
718
719        // Step 1: Get current table metadata
720        let metadata = self.table.metadata();
721        let mut next_field_id = metadata.last_column_id() + 1;
722        tracing::debug!("Starting schema change, next_field_id: {}", next_field_id);
723
724        // Step 2: Build new fields to add
725        let iceberg_create_table_arrow_convert = IcebergCreateTableArrowConvert::default();
726        let mut new_fields = Vec::new();
727
728        let mut drop_column_names = Vec::new();
729        match schema_change.op.as_ref() {
730            Some(risingwave_pb::stream_plan::sink_schema_change::Op::AddColumns(
731                add_columns_op,
732            )) => {
733                let add_columns = add_columns_op.fields.iter().map(Field::from).collect_vec();
734                for field in &add_columns {
735                    // Convert RisingWave Field to Arrow Field using IcebergCreateTableArrowConvert
736                    let arrow_field = iceberg_create_table_arrow_convert
737                        .to_arrow_field(&field.name, &field.data_type)
738                        .with_context(|| {
739                            format!("Failed to convert field '{}' to arrow", field.name)
740                        })
741                        .map_err(SinkError::Iceberg)?;
742
743                    // Convert Arrow DataType to Iceberg Type
744                    let iceberg_type = iceberg::arrow::arrow_type_to_type(arrow_field.data_type())
745                        .map_err(|err| {
746                            SinkError::Iceberg(
747                                anyhow!(err)
748                                    .context("Failed to convert Arrow type to Iceberg type"),
749                            )
750                        })?;
751
752                    // Create NestedField with the next available field ID
753                    let nested_field = Arc::new(NestedField::optional(
754                        next_field_id,
755                        &field.name,
756                        iceberg_type,
757                    ));
758
759                    new_fields.push(nested_field);
760                    tracing::info!("Prepared field '{}' with ID {}", field.name, next_field_id);
761                    next_field_id += 1;
762                }
763            }
764            Some(risingwave_pb::stream_plan::sink_schema_change::Op::DropColumns(
765                drop_columns_op,
766            )) => {
767                drop_column_names = drop_columns_op.column_names.clone();
768            }
769            _ => {
770                return Err(SinkError::Iceberg(anyhow!(
771                    "Unsupported sink schema change op in iceberg sink: {:?}",
772                    schema_change.op
773                )));
774            }
775        }
776
777        // Step 3: Create Transaction with UpdateSchemaAction
778        tracing::info!(
779            "Committing schema change to catalog for table {}",
780            self.table.identifier()
781        );
782
783        let txn = Transaction::new(&self.table);
784        let action_fields_added = new_fields.len();
785        let action = txn
786            .update_schema()
787            .add_fields(new_fields)
788            .drop_fields(drop_column_names.clone());
789
790        let updated_table = action
791            .apply(txn)
792            .context("Failed to apply schema update action")
793            .map_err(SinkError::Iceberg)?
794            .commit(self.catalog.as_ref())
795            .await
796            .context("Failed to commit table schema change")
797            .map_err(SinkError::Iceberg)?;
798
799        self.table = updated_table;
800
801        tracing::info!(
802            "Successfully committed schema change, added {} columns and dropped {} columns from iceberg table",
803            action_fields_added,
804            drop_column_names.len()
805        );
806
807        Ok(())
808    }
809
810    /// Check the number of snapshots on the given branch lineage since the last rewrite operation.
811    /// Returns the number of snapshots since the last rewrite.
812    fn count_snapshots_since_rewrite_in_metadata(metadata: &TableMetadata, branch: &str) -> usize {
813        // Start from the latest snapshot of the commit branch.
814        let mut snapshot_id = metadata
815            .snapshot_for_ref(branch)
816            .map(|snapshot| snapshot.snapshot_id());
817        let mut count = 0;
818
819        // Iterate through snapshots by parent lineage to find the last rewrite.
820        while let Some(current_snapshot_id) = snapshot_id {
821            let Some(snapshot) = metadata.snapshot_by_id(current_snapshot_id) else {
822                break;
823            };
824
825            // Check if this snapshot represents a rewrite operation.
826            if snapshot.summary().operation == Operation::Replace {
827                // Found a rewrite operation, stop counting.
828                break;
829            }
830
831            // Increment count for each snapshot that is not a rewrite.
832            count += 1;
833            snapshot_id = snapshot.parent_snapshot_id();
834        }
835
836        count
837    }
838
839    /// Returns the number of snapshots in the current commit branch since the last rewrite.
840    fn count_snapshots_since_rewrite(&self) -> usize {
841        let branch = commit_branch(self.config.r#type.as_str(), self.config.write_mode);
842        Self::count_snapshots_since_rewrite_in_metadata(self.table.metadata(), branch.as_str())
843    }
844
845    /// Wait until snapshot count since last rewrite is below the limit
846    async fn wait_for_snapshot_limit(&mut self) -> Result<()> {
847        if let Some(max_snapshots) = self.config.max_snapshots_num_before_compaction {
848            loop {
849                let current_count = self.count_snapshots_since_rewrite();
850
851                if current_count < max_snapshots {
852                    tracing::info!(
853                        "Snapshot count check passed: {} < {}",
854                        current_count,
855                        max_snapshots
856                    );
857                    break;
858                }
859
860                tracing::info!(
861                    "Snapshot count {} exceeds limit {}, waiting...",
862                    current_count,
863                    max_snapshots
864                );
865
866                self.notify_iceberg_compaction_scheduler(true);
867
868                // Wait for 30 seconds before checking again
869                tokio::time::sleep(Duration::from_secs(30)).await;
870
871                // Refresh table after the wait so the next check sees latest snapshots.
872                self.table = self.config.load_table().await?;
873            }
874        }
875        Ok(())
876    }
877}
878
879fn serialize_metadata(metadata: Vec<Vec<u8>>) -> Vec<u8> {
880    serde_json::to_vec(&metadata).unwrap()
881}
882
883fn deserialize_metadata(bytes: Vec<u8>) -> Vec<Vec<u8>> {
884    serde_json::from_slice(&bytes).unwrap()
885}
886
887#[cfg(test)]
888mod tests {
889    use std::collections::HashMap;
890
891    use iceberg::spec::{
892        FormatVersion, MAIN_BRANCH, NestedField, PrimitiveType, Schema, Snapshot,
893        SnapshotReference, SnapshotRetention, SortOrder, Summary, TableMetadataBuilder, Type,
894        UnboundPartitionSpec,
895    };
896    use risingwave_common::array::arrow::arrow_schema_iceberg::{
897        DataType as ArrowDataType, Field as ArrowField, FieldRef as ArrowFieldRef,
898        Fields as ArrowFields, Schema as ArrowSchema,
899    };
900
901    use super::*;
902
903    #[test]
904    fn test_schema_contains_same_fields_allows_binary_large_binary() {
905        let current_schema = ArrowSchema::new(vec![
906            ArrowField::new("k", ArrowDataType::Int32, true),
907            ArrowField::new("v", ArrowDataType::LargeBinary, true),
908        ]);
909        let expected_fields = vec![
910            ArrowField::new("k", ArrowDataType::Int32, true),
911            ArrowField::new("v", ArrowDataType::Binary, true),
912        ];
913
914        assert!(schema_contains_same_fields(
915            current_schema.fields(),
916            &expected_fields
917        ));
918    }
919
920    #[test]
921    fn test_schema_contains_same_fields_allows_nested_binary_large_binary() {
922        let current_schema = ArrowSchema::new(vec![
923            ArrowField::new(
924                "s",
925                ArrowDataType::Struct(ArrowFields::from(vec![ArrowField::new(
926                    "payload",
927                    ArrowDataType::LargeBinary,
928                    true,
929                )])),
930                true,
931            ),
932            ArrowField::new(
933                "l",
934                ArrowDataType::List(ArrowFieldRef::new(ArrowField::new_list_field(
935                    ArrowDataType::LargeBinary,
936                    true,
937                ))),
938                true,
939            ),
940            ArrowField::new_map(
941                "m",
942                "entries",
943                ArrowFieldRef::new(ArrowField::new("key", ArrowDataType::Utf8, false)),
944                ArrowFieldRef::new(ArrowField::new("value", ArrowDataType::LargeBinary, true)),
945                false,
946                true,
947            ),
948        ]);
949        let expected_fields = vec![
950            ArrowField::new(
951                "s",
952                ArrowDataType::Struct(ArrowFields::from(vec![ArrowField::new(
953                    "payload",
954                    ArrowDataType::Binary,
955                    true,
956                )])),
957                true,
958            ),
959            ArrowField::new(
960                "l",
961                ArrowDataType::List(ArrowFieldRef::new(ArrowField::new_list_field(
962                    ArrowDataType::Binary,
963                    true,
964                ))),
965                true,
966            ),
967            ArrowField::new_map(
968                "m",
969                "entries",
970                ArrowFieldRef::new(ArrowField::new("key", ArrowDataType::Utf8, false)),
971                ArrowFieldRef::new(ArrowField::new("value", ArrowDataType::Binary, true)),
972                false,
973                true,
974            ),
975        ];
976
977        assert!(schema_contains_same_fields(
978            current_schema.fields(),
979            &expected_fields
980        ));
981    }
982
983    #[test]
984    fn test_schema_contains_same_fields_allows_decimal_precision_delta() {
985        let current_schema = ArrowSchema::new(vec![ArrowField::new(
986            "d",
987            ArrowDataType::Decimal128(28, 10),
988            true,
989        )]);
990        let expected_fields = vec![ArrowField::new(
991            "d",
992            ArrowDataType::Decimal128(38, 10),
993            true,
994        )];
995
996        assert!(schema_contains_same_fields(
997            current_schema.fields(),
998            &expected_fields
999        ));
1000    }
1001
1002    #[test]
1003    fn test_schema_contains_same_fields_allows_decimal_precision_and_scale_delta() {
1004        let current_schema = ArrowSchema::new(vec![ArrowField::new(
1005            "d",
1006            ArrowDataType::Decimal128(38, 2),
1007            true,
1008        )]);
1009        let expected_fields = vec![ArrowField::new(
1010            "d",
1011            ArrowDataType::Decimal128(38, 10),
1012            true,
1013        )];
1014
1015        assert!(schema_contains_same_fields(
1016            current_schema.fields(),
1017            &expected_fields
1018        ));
1019    }
1020
1021    #[test]
1022    fn test_schema_contains_same_fields_rejects_length_mismatch() {
1023        let current_schema =
1024            ArrowSchema::new(vec![ArrowField::new("k", ArrowDataType::Int32, true)]);
1025        let expected_fields = vec![
1026            ArrowField::new("k", ArrowDataType::Int32, true),
1027            ArrowField::new("v", ArrowDataType::Utf8, true),
1028        ];
1029
1030        assert!(!schema_contains_same_fields(
1031            current_schema.fields(),
1032            &expected_fields
1033        ));
1034    }
1035
1036    #[test]
1037    fn test_schema_contains_same_fields_rejects_type_mismatch() {
1038        let current_schema =
1039            ArrowSchema::new(vec![ArrowField::new("v", ArrowDataType::Utf8, true)]);
1040        let expected_fields = vec![ArrowField::new("v", ArrowDataType::Int32, true)];
1041
1042        assert!(!schema_contains_same_fields(
1043            current_schema.fields(),
1044            &expected_fields
1045        ));
1046    }
1047
1048    #[test]
1049    fn test_schema_contains_same_fields_rejects_reused_duplicate_match() {
1050        let current_schema = ArrowSchema::new(vec![
1051            ArrowField::new("v", ArrowDataType::Int32, true),
1052            ArrowField::new("v", ArrowDataType::Utf8, true),
1053        ]);
1054        let expected_fields = vec![
1055            ArrowField::new("v", ArrowDataType::Int32, true),
1056            ArrowField::new("v", ArrowDataType::Int32, true),
1057        ];
1058
1059        assert!(!schema_contains_same_fields(
1060            current_schema.fields(),
1061            &expected_fields
1062        ));
1063    }
1064
1065    #[test]
1066    fn test_schema_contains_same_fields_rejects_nested_type_mismatch() {
1067        let current_schema = ArrowSchema::new(vec![ArrowField::new(
1068            "s",
1069            ArrowDataType::Struct(ArrowFields::from(vec![ArrowField::new(
1070                "payload",
1071                ArrowDataType::Utf8,
1072                true,
1073            )])),
1074            true,
1075        )]);
1076        let expected_fields = vec![ArrowField::new(
1077            "s",
1078            ArrowDataType::Struct(ArrowFields::from(vec![ArrowField::new(
1079                "payload",
1080                ArrowDataType::Int32,
1081                true,
1082            )])),
1083            true,
1084        )];
1085
1086        assert!(!schema_contains_same_fields(
1087            current_schema.fields(),
1088            &expected_fields
1089        ));
1090    }
1091
1092    #[test]
1093    fn test_count_snapshots_since_rewrite_in_metadata_ignores_other_branches() {
1094        let mut builder = TableMetadataBuilder::new(
1095            Schema::builder()
1096                .with_fields(vec![
1097                    NestedField::new(1, "id", Type::Primitive(PrimitiveType::Long), false).into(),
1098                ])
1099                .build()
1100                .unwrap(),
1101            UnboundPartitionSpec::builder().build(),
1102            SortOrder::unsorted_order(),
1103            "s3://warehouse/db/table".to_owned(),
1104            FormatVersion::V2,
1105            HashMap::new(),
1106        )
1107        .unwrap();
1108
1109        for (snapshot_id, parent_snapshot_id, operation) in [
1110            (1, None, Operation::Append),
1111            (2, Some(1), Operation::Append),
1112            (3, Some(2), Operation::Replace),
1113            (4, Some(3), Operation::Append),
1114            // Simulate the COW publish snapshot on main. It should not affect
1115            // the ingestion branch backlog count.
1116            (5, None, Operation::Overwrite),
1117        ] {
1118            builder = builder
1119                .add_snapshot(snapshot(snapshot_id, parent_snapshot_id, operation))
1120                .unwrap();
1121        }
1122
1123        let metadata = builder
1124            .set_ref(super::super::ICEBERG_COW_BRANCH, snapshot_ref(4))
1125            .unwrap()
1126            .set_ref(MAIN_BRANCH, snapshot_ref(5))
1127            .unwrap()
1128            .build()
1129            .unwrap()
1130            .metadata;
1131
1132        let count = IcebergSinkCommitter::count_snapshots_since_rewrite_in_metadata(
1133            &metadata,
1134            super::super::ICEBERG_COW_BRANCH,
1135        );
1136
1137        assert_eq!(count, 1);
1138    }
1139
1140    fn snapshot(
1141        snapshot_id: i64,
1142        parent_snapshot_id: Option<i64>,
1143        operation: Operation,
1144    ) -> Snapshot {
1145        Snapshot::builder()
1146            .with_snapshot_id(snapshot_id)
1147            .with_parent_snapshot_id(parent_snapshot_id)
1148            .with_sequence_number(snapshot_id)
1149            .with_timestamp_ms(snapshot_id)
1150            .with_manifest_list(format!("/snap-{snapshot_id}.avro"))
1151            .with_summary(Summary {
1152                operation,
1153                additional_properties: HashMap::new(),
1154            })
1155            .with_schema_id(0)
1156            .build()
1157    }
1158
1159    fn snapshot_ref(snapshot_id: i64) -> SnapshotReference {
1160        SnapshotReference::new(snapshot_id, SnapshotRetention::branch(None, None, None))
1161    }
1162}
risingwave_connector/sink/iceberg/commit.rs

risingwave_connector/sink/iceberg/
commit.rs
