Skip to main content

mz_storage/sink/
iceberg.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Iceberg sink implementation.
11//!
12//! This code renders a [`IcebergSinkConnection`] into a dataflow that writes
13//! data to an Iceberg table. `SinkRender::render_sink` hands the sink a stream
14//! of arrangement batches keyed on the sink key (the upstream arrangement's
15//! trace reader is already dropped, so the spine is free to compact as
16//! batches flow). A small `walk_sink_arrangement` operator consumes that
17//! stream and emits one `DiffPair` per `(key, timestamp)` update into the
18//! pipeline below.
19//!
20//! ```text
21//!        ┏━━━━━━━━━━━━━━┓
22//!        ┃   persist    ┃
23//!        ┃    source    ┃
24//!        ┗━━━━━━┯━━━━━━━┛
25//!               │ stream of arrangement batches (trace reader dropped)
26//!               │
27//!        ┏━━━━━━v━━━━━━━┓
28//!        ┃    walk      ┃
29//!        ┃ arrangement  ┃ yields individual DiffPairs per (key, timestamp)
30//!        ┗━━━━━━┯━━━━━━━┛
31//!               │ (Option<Row>, DiffPair<Row>) rows
32//!               │
33//!        ┏━━━━━━v━━━━━━━┓
34//!        ┃     mint     ┃ (single worker)
35//!        ┃    batch     ┃ loads/creates the Iceberg table,
36//!        ┃ descriptions ┃ determines resume upper
37//!        ┗━━━┯━━━━━┯━━━━┛
38//!            │     │ batch descriptions (broadcast)
39//!       rows │     ├─────────────────────────┐
40//!            │     │                         │
41//!        ┏━━━v━━━━━v━━━━┓    ╭─────────────╮ │
42//!        ┃    write     ┃───>│ S3 / object │ │
43//!        ┃  data files  ┃    │   storage   │ │
44//!        ┗━━━━━━┯━━━━━━━┛    ╰─────────────╯ │
45//!               │ file metadata              │
46//!               │                            │
47//!        ┏━━━━━━v━━━━━━━━━━━━━━━━━━━━━━━━━━━━v┓
48//!        ┃           commit to                ┃ (single worker)
49//!        ┃             iceberg                ┃
50//!        ┗━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━━━━━━┛
51//!                      │
52//!              ╭───────v───────╮
53//!              │ Iceberg table │
54//!              │  (snapshots)  │
55//!              ╰───────────────╯
56//! ```
57//! # Minting batch descriptions
58//! The "mint batch descriptions" operator is responsible for generating
59//! time-based batch boundaries that group writes into Iceberg snapshots.
60//! It maintains a sliding window of future batch descriptions so that
61//! writers can start processing data even while earlier batches are still being written.
62//! Knowing the batch boundaries ahead of time is important because we need to
63//! be able to make the claim that all data files written for a given batch
64//! include all data up to the upper `t` but not beyond it.
65//! This could be trivially achieved by waiting for all data to arrive up to a certain
66//! frontier, but that would prevent us from streaming writes out to object storage
67//! until the entire batch is complete, which would increase latency and reduce throughput.
68//!
69//! # Writing data files
70//! The "write data files" operator receives rows along with batch descriptions.
71//! It matches rows to batches by timestamp; if a batch description hasn't arrived yet,
72//! rows are stashed until it does. This allows batches to be minted ahead of data arrival.
73//! The operator uses an Iceberg `DeltaWriter` to write Parquet data files
74//! (and position delete files if necessary) to object storage.
75//! It outputs metadata about the written files along with their batch descriptions
76//! for the commit operator to consume.
77//!
78//! # Committing to Iceberg
79//! The "commit to iceberg" operator receives metadata about written data files
80//! along with their batch descriptions. It groups files by batch and creates
81//! Iceberg snapshots that include all files for each batch. It updates the Iceberg
82//! table's metadata to reflect the new snapshots, including updating the
83//! `mz-frontier` property to track progress.
84
85use std::cmp::Ordering;
86use std::collections::VecDeque;
87use std::convert::Infallible;
88use std::future::Future;
89use std::time::Instant;
90use std::{cell::RefCell, rc::Rc, sync::Arc};
91
92use anyhow::{Context, anyhow};
93use arrow::array::{ArrayRef, Int32Array, Int64Array, RecordBatch};
94use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
95use differential_dataflow::trace::BatchReader;
96use differential_dataflow::trace::implementations::ord_neu::OrdValBatch;
97use differential_dataflow::trace::implementations::{BatchContainer, Layout};
98use differential_dataflow::{Hashable, VecCollection};
99use futures::StreamExt;
100use iceberg::ErrorKind;
101use iceberg::arrow::{arrow_schema_to_schema, schema_to_arrow_schema};
102use iceberg::spec::{
103    DataFile, FormatVersion, Snapshot, StructType, read_data_files_from_avro,
104    write_data_files_to_avro,
105};
106use iceberg::spec::{Schema, SchemaRef};
107use iceberg::table::Table;
108use iceberg::transaction::{ApplyTransactionAction, Transaction};
109use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
110use iceberg::writer::base_writer::equality_delete_writer::{
111    EqualityDeleteFileWriterBuilder, EqualityDeleteWriterConfig,
112};
113use iceberg::writer::base_writer::position_delete_writer::{
114    PositionDeleteFileWriterBuilder, PositionDeleteWriterConfig,
115};
116use iceberg::writer::combined_writer::delta_writer::DeltaWriterBuilder;
117use iceberg::writer::file_writer::ParquetWriterBuilder;
118use iceberg::writer::file_writer::location_generator::{
119    DefaultFileNameGenerator, DefaultLocationGenerator,
120};
121use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
122use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
123use iceberg::{Catalog, NamespaceIdent, TableCreation, TableIdent};
124use itertools::Itertools;
125use mz_arrow_util::builder::{ARROW_EXTENSION_NAME_KEY, ArrowBuilder};
126use mz_interchange::avro::DiffPair;
127use mz_interchange::envelopes::for_each_diff_pair_async;
128use mz_ore::cast::CastFrom;
129use mz_ore::error::ErrorExt;
130use mz_ore::future::InTask;
131use mz_ore::result::ResultExt;
132use mz_ore::retry::{Retry, RetryResult};
133use mz_persist_client::Diagnostics;
134use mz_persist_client::write::WriteHandle;
135use mz_persist_types::codec_impls::UnitSchema;
136use mz_repr::{Diff, GlobalId, Row, Timestamp};
137use mz_row_spine::ArcBatch;
138use mz_storage_types::StorageDiff;
139use mz_storage_types::configuration::StorageConfiguration;
140use mz_storage_types::controller::CollectionMetadata;
141use mz_storage_types::errors::DataflowError;
142use mz_storage_types::sinks::{
143    IcebergSinkConnection, SinkEnvelope, StorageSinkDesc, iceberg_type_overrides,
144};
145use mz_storage_types::sources::SourceData;
146use mz_timely_util::antichain::AntichainExt;
147use mz_timely_util::builder_async::{Event, OperatorBuilder, PressOnDropButton};
148use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
149use parquet::file::properties::WriterProperties;
150use serde::{Deserialize, Serialize};
151use timely::PartialOrder;
152use timely::container::CapacityContainerBuilder;
153use timely::dataflow::StreamVec;
154use timely::dataflow::channels::pact::{Exchange, Pipeline};
155use timely::dataflow::operators::vec::{Broadcast, Map, ToStream};
156use timely::dataflow::operators::{CapabilitySet, Concatenate};
157use timely::progress::{Antichain, Timestamp as _};
158use tracing::debug;
159
160use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace};
161use crate::metrics::sink::iceberg::IcebergSinkMetrics;
162use crate::render::sinks::{PkViolationWarner, SinkBatchStream, SinkRender};
163use crate::statistics::SinkStatistics;
164use crate::storage_state::StorageState;
165
166/// Set the default capacity for the array builders inside the ArrowBuilder. This is the
167/// number of items each builder can hold before it needs to allocate more memory.
168const DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY: usize = 1024;
169/// Set the default buffer capacity for the string and binary array builders inside the
170/// ArrowBuilder. This is the number of bytes each builder can hold before it needs to allocate
171/// more memory.
172const DEFAULT_ARRAY_BUILDER_DATA_CAPACITY: usize = 1024;
173
174/// The prefix for Parquet files written by this sink.
175const PARQUET_FILE_PREFIX: &str = "mz_data";
176/// The number of batch descriptions to mint ahead of the observed frontier. This determines how
177/// many batches we have in-flight at any given time.
178const INITIAL_DESCRIPTIONS_TO_MINT: u64 = 3;
179
180/// Shared state produced by the async setup in [`write_data_files`] that both
181/// envelope handlers need to construct Parquet writers.
182struct WriterContext {
183    /// Arrow schema for data columns, with Materialize extension metadata merged in.
184    arrow_schema: Arc<ArrowSchema>,
185    /// Iceberg table schema, used to configure Parquet writers.
186    current_schema: Arc<Schema>,
187    /// File I/O for writing Parquet files to object storage.
188    file_io: iceberg::io::FileIO,
189    /// Generates file paths under the table's data directory.
190    location_generator: DefaultLocationGenerator,
191    /// Generates unique file names with a per-worker UUID suffix.
192    file_name_generator: DefaultFileNameGenerator,
193    writer_properties: WriterProperties,
194}
195
196/// Envelope-specific logic for writing Iceberg data files.
197trait EnvelopeHandler: Send {
198    /// Construct from the shared writer context after async setup completes.
199    fn new(
200        ctx: WriterContext,
201        connection: &IcebergSinkConnection,
202        materialize_arrow_schema: &Arc<ArrowSchema>,
203    ) -> anyhow::Result<Self>
204    where
205        Self: Sized;
206
207    /// Create an [`IcebergWriter`] for a new batch.
208    ///
209    /// `is_snapshot` is true for the initial "snapshot" batch (lower == as_of), which
210    /// contains all pre-existing data and can be very large. Implementations may use
211    /// this to disable memory-intensive optimisations like seen-rows deduplication.
212    async fn create_writer(&self, is_snapshot: bool) -> anyhow::Result<Box<dyn IcebergWriter>>;
213
214    fn row_to_batch(&self, diff_pair: DiffPair<Row>, ts: Timestamp) -> anyhow::Result<RecordBatch>;
215}
216
217struct UpsertEnvelopeHandler {
218    ctx: WriterContext,
219    /// Iceberg field IDs of the key columns, used for equality delete files.
220    equality_ids: Vec<i32>,
221    /// Iceberg schema for position delete files.
222    pos_schema: Arc<Schema>,
223    /// Iceberg schema for equality delete files (projected to key columns only).
224    eq_schema: Arc<Schema>,
225    /// Configuration for the equality delete writer (projected schema + column IDs).
226    eq_config: EqualityDeleteWriterConfig,
227    /// Arrow schema with an appended `__op` column that the
228    /// [`DeltaWriter`](iceberg::writer::combined_writer::delta_writer::DeltaWriter)
229    /// uses to distinguish inserts (+1) from deletes (-1).
230    schema_with_op: Arc<ArrowSchema>,
231}
232
233impl EnvelopeHandler for UpsertEnvelopeHandler {
234    fn new(
235        ctx: WriterContext,
236        connection: &IcebergSinkConnection,
237        materialize_arrow_schema: &Arc<ArrowSchema>,
238    ) -> anyhow::Result<Self> {
239        let Some((_, equality_indices)) = &connection.key_desc_and_indices else {
240            return Err(anyhow::anyhow!(
241                "Iceberg sink requires key columns for equality deletes"
242            ));
243        };
244
245        let equality_ids = equality_ids_for_indices(
246            ctx.current_schema.as_ref(),
247            materialize_arrow_schema.as_ref(),
248            equality_indices,
249        )?;
250
251        let pos_arrow_schema = PositionDeleteWriterConfig::arrow_schema();
252        let pos_schema = Arc::new(
253            arrow_schema_to_schema(&pos_arrow_schema)
254                .context("Failed to convert position delete Arrow schema to Iceberg schema")?,
255        );
256
257        let eq_config =
258            EqualityDeleteWriterConfig::new(equality_ids.clone(), Arc::clone(&ctx.current_schema))
259                .context("Failed to create EqualityDeleteWriterConfig")?;
260        let eq_schema = Arc::new(
261            arrow_schema_to_schema(eq_config.projected_arrow_schema_ref())
262                .context("Failed to convert equality delete Arrow schema to Iceberg schema")?,
263        );
264
265        let schema_with_op = Arc::new(build_schema_with_op_column(&ctx.arrow_schema));
266
267        Ok(Self {
268            ctx,
269            equality_ids,
270            pos_schema,
271            eq_schema,
272            eq_config,
273            schema_with_op,
274        })
275    }
276
277    async fn create_writer(&self, is_snapshot: bool) -> anyhow::Result<Box<dyn IcebergWriter>> {
278        let data_parquet_writer = ParquetWriterBuilder::new(
279            self.ctx.writer_properties.clone(),
280            Arc::clone(&self.ctx.current_schema),
281        )
282        .with_arrow_schema(Arc::clone(&self.ctx.arrow_schema))
283        .context("Arrow schema validation failed")?;
284        let data_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
285            data_parquet_writer,
286            Arc::clone(&self.ctx.current_schema),
287            self.ctx.file_io.clone(),
288            self.ctx.location_generator.clone(),
289            self.ctx.file_name_generator.clone(),
290        );
291        let data_writer_builder = DataFileWriterBuilder::new(data_rolling_writer);
292
293        let pos_config = PositionDeleteWriterConfig::new(None, 0, None);
294        let pos_parquet_writer = ParquetWriterBuilder::new(
295            self.ctx.writer_properties.clone(),
296            Arc::clone(&self.pos_schema),
297        );
298        let pos_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
299            pos_parquet_writer,
300            Arc::clone(&self.ctx.current_schema),
301            self.ctx.file_io.clone(),
302            self.ctx.location_generator.clone(),
303            self.ctx.file_name_generator.clone(),
304        );
305        let pos_delete_writer_builder =
306            PositionDeleteFileWriterBuilder::new(pos_rolling_writer, pos_config);
307
308        let eq_parquet_writer = ParquetWriterBuilder::new(
309            self.ctx.writer_properties.clone(),
310            Arc::clone(&self.eq_schema),
311        );
312        let eq_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
313            eq_parquet_writer,
314            Arc::clone(&self.ctx.current_schema),
315            self.ctx.file_io.clone(),
316            self.ctx.location_generator.clone(),
317            self.ctx.file_name_generator.clone(),
318        );
319        let eq_delete_writer_builder =
320            EqualityDeleteFileWriterBuilder::new(eq_rolling_writer, self.eq_config.clone());
321
322        let mut builder = DeltaWriterBuilder::new(
323            data_writer_builder,
324            pos_delete_writer_builder,
325            eq_delete_writer_builder,
326            self.equality_ids.clone(),
327        );
328
329        builder = if is_snapshot {
330            // Snapshot batches only produce inserts, so disable seen_rows tracking to save memory.
331            builder.with_max_seen_rows(0)
332        } else {
333            // For incremental batches, keep all "seen" rows. Do not evict any rows.
334            // The DeltaWriter issues an equality delete if we update (or delete) a row outside the "seen" cache.
335            // But equality deletes only apply to prior snapshots (lower sequence number).
336            //
337            // i.e. The DeltaWriter assumes that rows outside the "seen" cache come from prior snapshots.
338            //
339            // If we insert a row a=foo during this snapshot and then evict it from the "seen" cache,
340            // a subsequent update a=bar (also during this snapshot) will lead to:
341            //   1. Equality delete for a=foo (does nothing because a=foo is from this snapshot, not a prior snapshot)
342            //   2. Insert a=bar
343            // Because the deletion does nothing, we have a=foo and a=bar in the same snapshot.
344            builder.with_max_seen_rows(usize::MAX)
345        };
346
347        Ok(Box::new(
348            builder
349                .build(None)
350                .await
351                .context("Failed to create DeltaWriter")?,
352        ))
353    }
354
355    /// The `__op` column indicates whether each row is an insert (+1) or delete (-1),
356    /// which the DeltaWriter uses to generate the appropriate Iceberg data/delete files.
357    fn row_to_batch(
358        &self,
359        diff_pair: DiffPair<Row>,
360        _ts: Timestamp,
361    ) -> anyhow::Result<RecordBatch> {
362        let mut builder = ArrowBuilder::new_with_schema(
363            Arc::clone(&self.ctx.arrow_schema),
364            DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY,
365            DEFAULT_ARRAY_BUILDER_DATA_CAPACITY,
366        )
367        .context("Failed to create builder")?;
368
369        let mut op_values = Vec::new();
370
371        if let Some(before) = diff_pair.before {
372            builder
373                .add_row(&before)
374                .context("Failed to add delete row to builder")?;
375            op_values.push(-1i32);
376        }
377        if let Some(after) = diff_pair.after {
378            builder
379                .add_row(&after)
380                .context("Failed to add insert row to builder")?;
381            op_values.push(1i32);
382        }
383
384        let batch = builder
385            .to_record_batch()
386            .context("Failed to create record batch")?;
387
388        let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
389        columns.push(Arc::new(Int32Array::from(op_values)));
390
391        RecordBatch::try_new(Arc::clone(&self.schema_with_op), columns)
392            .context("Failed to create batch with op column")
393    }
394}
395
396struct AppendEnvelopeHandler {
397    ctx: WriterContext,
398    /// Arrow schema with only user columns (no `_mz_diff`/`_mz_timestamp`), used by
399    /// [`ArrowBuilder`] to serialize row data before the extra columns are appended.
400    user_schema_for_append: Arc<ArrowSchema>,
401}
402
403impl EnvelopeHandler for AppendEnvelopeHandler {
404    fn new(
405        ctx: WriterContext,
406        _connection: &IcebergSinkConnection,
407        _materialize_arrow_schema: &Arc<ArrowSchema>,
408    ) -> anyhow::Result<Self> {
409        // arrow_schema already includes _mz_diff + _mz_timestamp (added in render_sink); strip
410        // the last two fields so ArrowBuilder only processes the user columns.
411        let n = ctx.arrow_schema.fields().len().saturating_sub(2);
412        let user_schema_for_append =
413            Arc::new(ArrowSchema::new(ctx.arrow_schema.fields()[..n].to_vec()));
414
415        Ok(Self {
416            ctx,
417            user_schema_for_append,
418        })
419    }
420
421    async fn create_writer(&self, _is_snapshot: bool) -> anyhow::Result<Box<dyn IcebergWriter>> {
422        let data_parquet_writer = ParquetWriterBuilder::new(
423            self.ctx.writer_properties.clone(),
424            Arc::clone(&self.ctx.current_schema),
425        )
426        .with_arrow_schema(Arc::clone(&self.ctx.arrow_schema))
427        .context("Arrow schema validation failed")?;
428        let data_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
429            data_parquet_writer,
430            Arc::clone(&self.ctx.current_schema),
431            self.ctx.file_io.clone(),
432            self.ctx.location_generator.clone(),
433            self.ctx.file_name_generator.clone(),
434        );
435        Ok(Box::new(
436            DataFileWriterBuilder::new(data_rolling_writer)
437                .build(None)
438                .await
439                .context("Failed to create DataFileWriter")?,
440        ))
441    }
442
443    /// Every change is written as a plain data row: the `before` half (if present) gets
444    /// `_mz_diff = -1` and the `after` half gets `_mz_diff = +1`. Both carry the same `_mz_timestamp`.
445    fn row_to_batch(&self, diff_pair: DiffPair<Row>, ts: Timestamp) -> anyhow::Result<RecordBatch> {
446        let mut builder = ArrowBuilder::new_with_schema(
447            Arc::clone(&self.user_schema_for_append),
448            DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY,
449            DEFAULT_ARRAY_BUILDER_DATA_CAPACITY,
450        )
451        .context("Failed to create builder")?;
452
453        let mut diff_values: Vec<i32> = Vec::new();
454        let ts_i64 = i64::try_from(u64::from(ts)).unwrap_or(i64::MAX);
455
456        if let Some(before) = diff_pair.before {
457            builder
458                .add_row(&before)
459                .context("Failed to add before row to builder")?;
460            diff_values.push(-1i32);
461        }
462        if let Some(after) = diff_pair.after {
463            builder
464                .add_row(&after)
465                .context("Failed to add after row to builder")?;
466            diff_values.push(1i32);
467        }
468
469        let n = diff_values.len();
470        let batch = builder
471            .to_record_batch()
472            .context("Failed to create record batch")?;
473
474        let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
475        columns.push(Arc::new(Int32Array::from(diff_values)));
476        columns.push(Arc::new(Int64Array::from(vec![ts_i64; n])));
477
478        RecordBatch::try_new(Arc::clone(&self.ctx.arrow_schema), columns)
479            .context("Failed to create append record batch")
480    }
481}
482
483/// Add Parquet field IDs to an Arrow schema. Iceberg requires field IDs in the
484/// Parquet metadata for schema evolution tracking. Field IDs are assigned
485/// recursively to all nested fields (structs, lists, maps) using a depth-first,
486/// pre-order traversal.
487fn add_field_ids_to_arrow_schema(schema: ArrowSchema) -> ArrowSchema {
488    let mut next_field_id = 1i32;
489    let fields: Vec<Field> = schema
490        .fields()
491        .iter()
492        .map(|field| add_field_ids_recursive(field, &mut next_field_id))
493        .collect();
494    ArrowSchema::new(fields).with_metadata(schema.metadata().clone())
495}
496
497/// Recursively add field IDs to a field and all its nested children.
498fn add_field_ids_recursive(field: &Field, next_id: &mut i32) -> Field {
499    let current_id = *next_id;
500    *next_id += 1;
501
502    let mut metadata = field.metadata().clone();
503    metadata.insert(
504        PARQUET_FIELD_ID_META_KEY.to_string(),
505        current_id.to_string(),
506    );
507
508    let new_data_type = add_field_ids_to_datatype(field.data_type(), next_id);
509
510    Field::new(field.name(), new_data_type, field.is_nullable()).with_metadata(metadata)
511}
512
513/// Add field IDs to nested fields within a DataType.
514fn add_field_ids_to_datatype(data_type: &DataType, next_id: &mut i32) -> DataType {
515    match data_type {
516        DataType::Struct(fields) => {
517            let new_fields: Vec<Field> = fields
518                .iter()
519                .map(|f| add_field_ids_recursive(f, next_id))
520                .collect();
521            DataType::Struct(new_fields.into())
522        }
523        DataType::List(element_field) => {
524            let new_element = add_field_ids_recursive(element_field, next_id);
525            DataType::List(Arc::new(new_element))
526        }
527        DataType::LargeList(element_field) => {
528            let new_element = add_field_ids_recursive(element_field, next_id);
529            DataType::LargeList(Arc::new(new_element))
530        }
531        DataType::Map(entries_field, sorted) => {
532            let new_entries = add_field_ids_recursive(entries_field, next_id);
533            DataType::Map(Arc::new(new_entries), *sorted)
534        }
535        _ => data_type.clone(),
536    }
537}
538
539/// Merge Materialize extension metadata into Iceberg's Arrow schema.
540/// This uses Iceberg's data types (e.g. Utf8) and field IDs while preserving
541/// Materialize's extension names for ArrowBuilder compatibility.
542/// Handles nested types (structs, lists, maps) recursively.
543fn merge_materialize_metadata_into_iceberg_schema(
544    materialize_arrow_schema: &ArrowSchema,
545    iceberg_schema: &Schema,
546) -> anyhow::Result<ArrowSchema> {
547    // First, convert Iceberg schema to Arrow (this gives us the correct data types)
548    let iceberg_arrow_schema = schema_to_arrow_schema(iceberg_schema)
549        .context("Failed to convert Iceberg schema to Arrow schema")?;
550
551    // Now merge in the Materialize extension metadata
552    let fields: Vec<Field> = iceberg_arrow_schema
553        .fields()
554        .iter()
555        .map(|iceberg_field| {
556            // Find the corresponding Materialize field by name to get extension metadata
557            let mz_field = materialize_arrow_schema
558                .field_with_name(iceberg_field.name())
559                .with_context(|| {
560                    format!(
561                        "Field '{}' not found in Materialize schema",
562                        iceberg_field.name()
563                    )
564                })?;
565
566            merge_field_metadata_recursive(iceberg_field, Some(mz_field))
567        })
568        .collect::<anyhow::Result<Vec<_>>>()?;
569
570    Ok(ArrowSchema::new(fields).with_metadata(iceberg_arrow_schema.metadata().clone()))
571}
572
573/// Recursively merge Materialize extension metadata into an Iceberg field.
574fn merge_field_metadata_recursive(
575    iceberg_field: &Field,
576    mz_field: Option<&Field>,
577) -> anyhow::Result<Field> {
578    // Start with Iceberg field's metadata (which includes field IDs)
579    let mut metadata = iceberg_field.metadata().clone();
580
581    // Add Materialize extension name if available
582    if let Some(mz_f) = mz_field {
583        if let Some(extension_name) = mz_f.metadata().get(ARROW_EXTENSION_NAME_KEY) {
584            metadata.insert(ARROW_EXTENSION_NAME_KEY.to_string(), extension_name.clone());
585        }
586    }
587
588    // Recursively process nested types
589    let new_data_type = match iceberg_field.data_type() {
590        DataType::Struct(iceberg_fields) => {
591            let mz_struct_fields = match mz_field {
592                Some(f) => match f.data_type() {
593                    DataType::Struct(fields) => Some(fields),
594                    other => anyhow::bail!(
595                        "Type mismatch for field '{}': Iceberg schema has Struct, but Materialize schema has {:?}",
596                        iceberg_field.name(),
597                        other
598                    ),
599                },
600                None => None,
601            };
602
603            let new_fields: Vec<Field> = iceberg_fields
604                .iter()
605                .map(|iceberg_inner| {
606                    let mz_inner = mz_struct_fields.and_then(|fields| {
607                        fields.iter().find(|f| f.name() == iceberg_inner.name())
608                    });
609                    merge_field_metadata_recursive(iceberg_inner, mz_inner.map(|f| f.as_ref()))
610                })
611                .collect::<anyhow::Result<Vec<_>>>()?;
612
613            DataType::Struct(new_fields.into())
614        }
615        DataType::List(iceberg_element) => {
616            let mz_element = match mz_field {
617                Some(f) => match f.data_type() {
618                    DataType::List(element) => Some(element.as_ref()),
619                    other => anyhow::bail!(
620                        "Type mismatch for field '{}': Iceberg schema has List, but Materialize schema has {:?}",
621                        iceberg_field.name(),
622                        other
623                    ),
624                },
625                None => None,
626            };
627            let new_element = merge_field_metadata_recursive(iceberg_element, mz_element)?;
628            DataType::List(Arc::new(new_element))
629        }
630        DataType::LargeList(iceberg_element) => {
631            let mz_element = match mz_field {
632                Some(f) => match f.data_type() {
633                    DataType::LargeList(element) => Some(element.as_ref()),
634                    other => anyhow::bail!(
635                        "Type mismatch for field '{}': Iceberg schema has LargeList, but Materialize schema has {:?}",
636                        iceberg_field.name(),
637                        other
638                    ),
639                },
640                None => None,
641            };
642            let new_element = merge_field_metadata_recursive(iceberg_element, mz_element)?;
643            DataType::LargeList(Arc::new(new_element))
644        }
645        DataType::Map(iceberg_entries, sorted) => {
646            let mz_entries = match mz_field {
647                Some(f) => match f.data_type() {
648                    DataType::Map(entries, _) => Some(entries.as_ref()),
649                    other => anyhow::bail!(
650                        "Type mismatch for field '{}': Iceberg schema has Map, but Materialize schema has {:?}",
651                        iceberg_field.name(),
652                        other
653                    ),
654                },
655                None => None,
656            };
657            // The Iceberg arrow representation names map fields differently from
658            // Materialize (`key_value`/`key`/`value` vs `entries`/`keys`/`values`),
659            // so name-based matching on the entries struct would drop the value
660            // field's extension metadata. Merge the entries struct positionally.
661            let new_entries = match mz_entries {
662                Some(mz_entries) => merge_map_entries_metadata(iceberg_entries, mz_entries)?,
663                None => iceberg_entries.as_ref().clone(),
664            };
665            DataType::Map(Arc::new(new_entries), *sorted)
666        }
667        other => other.clone(),
668    };
669
670    Ok(Field::new(
671        iceberg_field.name(),
672        new_data_type,
673        iceberg_field.is_nullable(),
674    )
675    .with_metadata(metadata))
676}
677
678/// Merge metadata into a Map's entries struct, matching key/value positionally.
679///
680/// Iceberg's arrow representation names map fields `key_value`/`key`/`value`,
681/// while Materialize uses `entries`/`keys`/`values`. Name-based matching would
682/// drop the materialize extension metadata for the value field, which then
683/// causes `ArrowBuilder` to fail with "Field 'value' missing extension metadata".
684///
685/// Positional matching is safe because the Arrow spec defines Map structurally,
686/// not by field name: `List<entries: Struct<key: K, value: V>>` with exactly
687/// two struct children — key first, value second — and the names are only
688/// conventional. See `Map` in apache/arrow `format/Schema.fbs`:
689/// <https://github.com/apache/arrow/blob/main/format/Schema.fbs> — "The names
690/// of the child fields may be respectively 'entries', 'key', and 'value', but
691/// this is not enforced."
692///
693/// Future cleanup: we could instead align Materialize's arrow map field names
694/// with the Parquet/Iceberg convention (`key_value`/`key`/`value`) in
695/// `mz_arrow_util::builder::scalar_to_arrow_datatype_impl` and drop this
696/// positional helper. That would also affect `COPY TO S3 ... FORMAT = 'parquet'`
697/// output schemas, so we'd need to confirm no downstream consumers depend on
698/// the current `entries`/`keys`/`values` names before flipping.
699fn merge_map_entries_metadata(
700    iceberg_entries: &Field,
701    mz_entries: &Field,
702) -> anyhow::Result<Field> {
703    let mut metadata = iceberg_entries.metadata().clone();
704    if let Some(extension_name) = mz_entries.metadata().get(ARROW_EXTENSION_NAME_KEY) {
705        metadata.insert(ARROW_EXTENSION_NAME_KEY.to_string(), extension_name.clone());
706    }
707
708    let iceberg_fields = match iceberg_entries.data_type() {
709        DataType::Struct(fields) => fields,
710        other => anyhow::bail!(
711            "Iceberg map entries field '{}' is not a Struct: {:?}",
712            iceberg_entries.name(),
713            other
714        ),
715    };
716    let mz_fields = match mz_entries.data_type() {
717        DataType::Struct(fields) => fields,
718        other => anyhow::bail!(
719            "Materialize map entries field '{}' is not a Struct: {:?}",
720            mz_entries.name(),
721            other
722        ),
723    };
724
725    let new_fields: Vec<Field> = iceberg_fields
726        .iter()
727        .enumerate()
728        .map(|(idx, iceberg_inner)| {
729            let mz_inner = mz_fields.get(idx).map(|f| f.as_ref());
730            merge_field_metadata_recursive(iceberg_inner, mz_inner)
731        })
732        .collect::<anyhow::Result<Vec<_>>>()?;
733
734    Ok(Field::new(
735        iceberg_entries.name(),
736        DataType::Struct(new_fields.into()),
737        iceberg_entries.is_nullable(),
738    )
739    .with_metadata(metadata))
740}
741
742async fn reload_table(
743    catalog: &dyn Catalog,
744    namespace: String,
745    table_name: String,
746    current_table: Table,
747) -> anyhow::Result<Table> {
748    let namespace_ident = NamespaceIdent::new(namespace.clone());
749    let table_ident = TableIdent::new(namespace_ident, table_name.clone());
750    let current_schema = current_table.metadata().current_schema_id();
751    let current_partition_spec = current_table.metadata().default_partition_spec_id();
752
753    match catalog.load_table(&table_ident).await {
754        Ok(table) => {
755            let reloaded_schema = table.metadata().current_schema_id();
756            let reloaded_partition_spec = table.metadata().default_partition_spec_id();
757            if reloaded_schema != current_schema {
758                return Err(anyhow::anyhow!(
759                    "Iceberg table '{}' schema changed during operation but schema evolution isn't supported, expected schema ID {}, got {}",
760                    table_name,
761                    current_schema,
762                    reloaded_schema
763                ));
764            }
765
766            if reloaded_partition_spec != current_partition_spec {
767                return Err(anyhow::anyhow!(
768                    "Iceberg table '{}' partition spec changed during operation but partition spec evolution isn't supported, expected partition spec ID {}, got {}",
769                    table_name,
770                    current_partition_spec,
771                    reloaded_partition_spec
772                ));
773            }
774
775            Ok(table)
776        }
777        Err(err) => Err(err).context("Failed to reload Iceberg table"),
778    }
779}
780
781/// Attempt a single commit of a batch of data files to an Iceberg table.
782/// On conflict or failure, reloads the table and returns a retryable error.
783/// On success, returns the updated table state.
784async fn try_commit_batch(
785    mut table: Table,
786    snapshot_properties: Vec<(String, String)>,
787    data_files: Vec<DataFile>,
788    delete_files: Vec<DataFile>,
789    catalog: &dyn Catalog,
790    conn_namespace: &str,
791    conn_table: &str,
792    sink_version: u64,
793    frontier: &Antichain<Timestamp>,
794    batch_lower: &Antichain<Timestamp>,
795    batch_upper: &Antichain<Timestamp>,
796    metrics: &IcebergSinkMetrics,
797) -> (Table, RetryResult<(), anyhow::Error>) {
798    let tx = Transaction::new(&table);
799    let mut action = tx
800        .row_delta()
801        .set_snapshot_properties(snapshot_properties.into_iter().collect())
802        .with_check_duplicate(false);
803
804    if !data_files.is_empty() || !delete_files.is_empty() {
805        action = action
806            .add_data_files(data_files)
807            .add_delete_files(delete_files);
808    }
809
810    let tx = match action
811        .apply(tx)
812        .context("Failed to apply data file addition to iceberg table transaction")
813    {
814        Ok(tx) => tx,
815        Err(e) => {
816            match reload_table(
817                catalog,
818                conn_namespace.to_string(),
819                conn_table.to_string(),
820                table.clone(),
821            )
822            .await
823            {
824                Ok(reloaded) => table = reloaded,
825                Err(reload_err) => {
826                    return (table, RetryResult::RetryableErr(anyhow!(reload_err)));
827                }
828            }
829            return (
830                table,
831                RetryResult::RetryableErr(anyhow!(
832                    "Failed to apply data file addition to iceberg table transaction: {}",
833                    e
834                )),
835            );
836        }
837    };
838
839    let new_table = tx.commit(catalog).await;
840    match new_table {
841        Err(e) if matches!(e.kind(), ErrorKind::CatalogCommitConflicts) => {
842            metrics.commit_conflicts.inc();
843            match reload_table(
844                catalog,
845                conn_namespace.to_string(),
846                conn_table.to_string(),
847                table.clone(),
848            )
849            .await
850            {
851                Ok(reloaded) => table = reloaded,
852                Err(e) => {
853                    return (table, RetryResult::RetryableErr(anyhow!(e)));
854                }
855            };
856
857            let mut snapshots: Vec<_> = table.metadata().snapshots().cloned().collect();
858            let last = retrieve_upper_from_snapshots(&mut snapshots);
859            let last = match last {
860                Ok(val) => val,
861                Err(e) => {
862                    return (table, RetryResult::RetryableErr(anyhow!(e)));
863                }
864            };
865
866            // Check if another writer has advanced the frontier beyond ours (fencing check)
867            if let Some((last_frontier, last_version)) = last {
868                if last_version > sink_version {
869                    return (
870                        table,
871                        RetryResult::FatalErr(anyhow!(
872                            "Iceberg table '{}' has been modified by another writer \
873                             with version {}. Current sink version: {}. \
874                             Frontiers may be out of sync, aborting to avoid data loss.",
875                            conn_table,
876                            last_version,
877                            sink_version,
878                        )),
879                    );
880                }
881                if PartialOrder::less_equal(frontier, &last_frontier) {
882                    return (
883                        table,
884                        RetryResult::FatalErr(anyhow!(
885                            "Iceberg table '{}' has been modified by another writer. \
886                             Current frontier: {:?}, last frontier: {:?}.",
887                            conn_table,
888                            frontier,
889                            last_frontier,
890                        )),
891                    );
892                }
893            }
894
895            (
896                table,
897                RetryResult::RetryableErr(anyhow!(
898                    "Commit conflict detected when committing batch [{}, {}) \
899                     to Iceberg table '{}.{}'. Retrying...",
900                    batch_lower.pretty(),
901                    batch_upper.pretty(),
902                    conn_namespace,
903                    conn_table
904                )),
905            )
906        }
907        Err(e) => {
908            metrics.commit_failures.inc();
909            (table, RetryResult::RetryableErr(anyhow!(e)))
910        }
911        Ok(new_table) => (new_table, RetryResult::Ok(())),
912    }
913}
914
915/// Load an existing Iceberg table or create it if it doesn't exist.
916async fn load_or_create_table(
917    catalog: &dyn Catalog,
918    namespace: String,
919    table_name: String,
920    schema: &Schema,
921) -> anyhow::Result<iceberg::table::Table> {
922    let namespace_ident = NamespaceIdent::new(namespace.clone());
923    let table_ident = TableIdent::new(namespace_ident.clone(), table_name.clone());
924
925    // Try to load the table first
926    match catalog.load_table(&table_ident).await {
927        Ok(table) => {
928            // Table exists, return it
929            // TODO: Add proper schema evolution/validation to ensure compatibility
930            Ok(table)
931        }
932        Err(err) => {
933            if matches!(err.kind(), ErrorKind::TableNotFound { .. })
934                || err
935                    .message()
936                    .contains("Tried to load a table that does not exist")
937            {
938                // Table doesn't exist, create it
939                // Note: location is not specified, letting the catalog determine the default location
940                // based on its warehouse configuration
941                let table_creation = TableCreation::builder()
942                    .name(table_name.clone())
943                    .schema(schema.clone())
944                    // Use unpartitioned spec by default
945                    // TODO: Consider making partition spec configurable
946                    // .partition_spec(UnboundPartitionSpec::builder().build())
947                    .build();
948
949                catalog
950                    .create_table(&namespace_ident, table_creation)
951                    .await
952                    .with_context(|| {
953                        format!(
954                            "Failed to create Iceberg table '{}' in namespace '{}'",
955                            table_name, namespace
956                        )
957                    })
958            } else {
959                // Some other error occurred
960                Err(err).context("Failed to load Iceberg table")
961            }
962        }
963    }
964}
965
966/// Find the most recent Materialize frontier from Iceberg snapshots.
967/// We store the frontier in snapshot metadata to track where we left off after restarts.
968/// Snapshots with operation="replace" (compactions) don't have our metadata and are skipped.
969/// The input slice will be sorted by sequence number in descending order.
970fn retrieve_upper_from_snapshots(
971    snapshots: &mut [Arc<Snapshot>],
972) -> anyhow::Result<Option<(Antichain<Timestamp>, u64)>> {
973    snapshots.sort_by(|a, b| Ord::cmp(&b.sequence_number(), &a.sequence_number()));
974
975    for snapshot in snapshots {
976        let props = &snapshot.summary().additional_properties;
977        if let (Some(frontier_json), Some(sink_version_str)) =
978            (props.get("mz-frontier"), props.get("mz-sink-version"))
979        {
980            let frontier: Vec<Timestamp> = serde_json::from_str(frontier_json)
981                .context("Failed to deserialize frontier from snapshot properties")?;
982            let frontier = Antichain::from_iter(frontier);
983
984            let sink_version = sink_version_str
985                .parse::<u64>()
986                .context("Failed to parse mz-sink-version from snapshot properties")?;
987
988            return Ok(Some((frontier, sink_version)));
989        }
990        if snapshot.summary().operation.as_str() != "replace" {
991            // This is a bad heuristic, but we have no real other way to identify compactions
992            // right now other than assume they will be the only operation writing "replace" operations.
993            // That means if we find a snapshot with some other operation, but no mz-frontier, we are in an
994            // inconsistent state and have to error out.
995            anyhow::bail!(
996                "Iceberg table is in an inconsistent state: snapshot {} has operation '{}' but is missing 'mz-frontier' property. Schema or partition spec evolution is not supported.",
997                snapshot.snapshot_id(),
998                snapshot.summary().operation.as_str(),
999            );
1000        }
1001    }
1002
1003    Ok(None)
1004}
1005
1006/// Convert a Materialize RelationDesc into Arrow and Iceberg schemas.
1007///
1008/// Returns a tuple of:
1009/// - The Arrow schema (with field IDs and Iceberg-compatible types) for writing Parquet files
1010/// - The Iceberg schema for table creation/validation
1011///
1012/// Iceberg doesn't support unsigned integer types, so we use `iceberg_type_overrides`
1013/// to map them to compatible types (e.g., UInt64 -> Decimal128(20,0)). The ArrowBuilder
1014/// handles the cross-type conversion (Datum::UInt64 -> Decimal128Builder) automatically.
1015fn relation_desc_to_iceberg_schema(
1016    desc: &mz_repr::RelationDesc,
1017) -> anyhow::Result<(ArrowSchema, SchemaRef)> {
1018    let arrow_schema =
1019        mz_arrow_util::builder::desc_to_schema_with_overrides(desc, iceberg_type_overrides)
1020            .context("Failed to convert RelationDesc to Iceberg-compatible Arrow schema")?;
1021
1022    let arrow_schema_with_ids = add_field_ids_to_arrow_schema(arrow_schema);
1023
1024    let iceberg_schema = arrow_schema_to_schema(&arrow_schema_with_ids)
1025        .context("Failed to convert Arrow schema to Iceberg schema")?;
1026
1027    Ok((arrow_schema_with_ids, Arc::new(iceberg_schema)))
1028}
1029
1030/// Resolve Materialize key column indexes to Iceberg top-level field IDs.
1031///
1032/// Iceberg field IDs are assigned recursively, so a top-level column's field ID
1033/// is not necessarily `column_index + 1` once nested fields are present.
1034fn equality_ids_for_indices(
1035    current_schema: &Schema,
1036    materialize_arrow_schema: &ArrowSchema,
1037    equality_indices: &[usize],
1038) -> anyhow::Result<Vec<i32>> {
1039    let top_level_fields = current_schema.as_struct();
1040
1041    equality_indices
1042        .iter()
1043        .map(|index| {
1044            let mz_field = materialize_arrow_schema
1045                .fields()
1046                .get(*index)
1047                .with_context(|| format!("Equality delete key index {index} is out of bounds"))?;
1048            let field_name = mz_field.name();
1049            let iceberg_field = top_level_fields
1050                .field_by_name(field_name)
1051                .with_context(|| {
1052                    format!(
1053                        "Equality delete key column '{}' not found in Iceberg table schema",
1054                        field_name
1055                    )
1056                })?;
1057            Ok(iceberg_field.id)
1058        })
1059        .collect()
1060}
1061
1062/// Build a new Arrow schema by adding an __op column to the existing schema.
1063fn build_schema_with_op_column(schema: &ArrowSchema) -> ArrowSchema {
1064    let mut fields: Vec<Arc<Field>> = schema.fields().iter().cloned().collect();
1065    fields.push(Arc::new(Field::new("__op", DataType::Int32, false)));
1066    ArrowSchema::new(fields)
1067}
1068
1069/// Build a new Arrow schema by appending `_mz_diff` (Int32) and `_mz_timestamp` (Int64) columns.
1070/// These are user-visible Iceberg columns written in append mode. Parquet field IDs are
1071/// assigned sequentially after the existing maximum field ID so the extended schema can
1072/// be converted to a valid Iceberg schema via `arrow_schema_to_schema`.
1073#[allow(clippy::disallowed_types)]
1074fn build_schema_with_append_columns(schema: &ArrowSchema) -> ArrowSchema {
1075    use mz_storage_types::sinks::{ICEBERG_APPEND_DIFF_COLUMN, ICEBERG_APPEND_TIMESTAMP_COLUMN};
1076    let mut fields: Vec<Arc<Field>> = schema.fields().iter().cloned().collect();
1077    fields.push(Arc::new(Field::new(
1078        ICEBERG_APPEND_DIFF_COLUMN,
1079        DataType::Int32,
1080        false,
1081    )));
1082    fields.push(Arc::new(Field::new(
1083        ICEBERG_APPEND_TIMESTAMP_COLUMN,
1084        DataType::Int64,
1085        false,
1086    )));
1087
1088    add_field_ids_to_arrow_schema(ArrowSchema::new(fields).with_metadata(schema.metadata().clone()))
1089}
1090
1091/// Generate time-based batch boundaries for grouping writes into Iceberg snapshots.
1092/// Batches are minted with configurable windows to balance write efficiency with latency.
1093/// We maintain a sliding window of future batch descriptions so writers can start
1094/// processing data even while earlier batches are still being written.
1095fn mint_batch_descriptions<'scope>(
1096    name: String,
1097    sink_id: GlobalId,
1098    input: SinkBatchStream<'scope>,
1099    sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
1100    connection: IcebergSinkConnection,
1101    storage_configuration: StorageConfiguration,
1102    initial_schema: SchemaRef,
1103) -> (
1104    StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
1105    StreamVec<'scope, Timestamp, Infallible>,
1106    StreamVec<'scope, Timestamp, HealthStatusMessage>,
1107    PressOnDropButton,
1108) {
1109    let scope = input.scope();
1110    let name_for_error = name.clone();
1111    let name_for_logging = name.clone();
1112    let mut builder = OperatorBuilder::new(name, scope.clone());
1113    let sink_version = sink.version;
1114
1115    let hashed_id = sink_id.hashed();
1116    let is_active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
1117    let (_, table_ready_stream) = builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
1118    let (batch_desc_output, batch_desc_stream) =
1119        builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
1120    let mut input = builder.new_input_for(input, Pipeline, &batch_desc_output);
1121
1122    let as_of = sink.as_of.clone();
1123    let commit_interval = sink
1124        .commit_interval
1125        .expect("the planner should have enforced this")
1126        .clone();
1127
1128    let (button, errors): (_, StreamVec<'scope, Timestamp, Rc<anyhow::Error>>) =
1129        builder.build_fallible(move |caps| {
1130        Box::pin(async move {
1131            let [table_ready_capset, capset]: &mut [_; 2] = caps.try_into().unwrap();
1132
1133            if !is_active_worker {
1134                // Only the active worker mints batch descriptions.
1135                return Ok(());
1136            }
1137
1138            let catalog = connection
1139                .catalog_connection
1140                .connect(&storage_configuration, InTask::Yes)
1141                .await
1142                .with_context(|| {
1143                    format!(
1144                        "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
1145                        connection.catalog_connection.uri, connection.namespace, connection.table
1146                    )
1147                })?;
1148
1149            let table = load_or_create_table(
1150                catalog.as_ref(),
1151                connection.namespace.clone(),
1152                connection.table.clone(),
1153                initial_schema.as_ref(),
1154            )
1155            .await?;
1156            debug!(
1157                ?sink_id,
1158                %name_for_logging,
1159                namespace = %connection.namespace,
1160                table = %connection.table,
1161                "iceberg mint loaded/created table"
1162            );
1163
1164            *table_ready_capset = CapabilitySet::new();
1165
1166            let mut snapshots: Vec<_> = table.metadata().snapshots().cloned().collect();
1167            let resume = retrieve_upper_from_snapshots(&mut snapshots)?;
1168            let (resume_upper, resume_version) = match resume {
1169                Some((f, v)) => (f, v),
1170                None => (Antichain::from_elem(Timestamp::minimum()), 0),
1171            };
1172            debug!(
1173                ?sink_id,
1174                %name_for_logging,
1175                resume_upper = %resume_upper.pretty(),
1176                resume_version,
1177                as_of = %as_of.pretty(),
1178                "iceberg mint resume position loaded"
1179            );
1180
1181            // The input has overcompacted if
1182            let overcompacted =
1183                // ..we have made some progress in the past
1184                *resume_upper != [Timestamp::minimum()] &&
1185                // ..but the since frontier is now beyond that
1186                PartialOrder::less_than(&resume_upper, &as_of);
1187
1188            if overcompacted {
1189                let err = format!(
1190                    "{name_for_error}: input compacted past resume upper: as_of {}, resume_upper: {}",
1191                    as_of.pretty(),
1192                    resume_upper.pretty()
1193                );
1194                // This would normally be an assertion but because it can happen after a
1195                // Materialize backup/restore we log an error so that it appears on Sentry but
1196                // leaves the rest of the objects in the cluster unaffected.
1197                return Err(anyhow::anyhow!("{err}"));
1198            };
1199
1200            if resume_version > sink_version {
1201                anyhow::bail!("Fenced off by newer sink version: resume_version {}, sink_version {}", resume_version, sink_version);
1202            }
1203
1204            let mut initialized = false;
1205            let mut observed_frontier;
1206            // Track minted batches to maintain a sliding window of open batch descriptions.
1207            // This is needed to know when to retire old batches and mint new ones.
1208            // It's "sortedness" is derived from the monotonicity of batch descriptions,
1209            // and the fact that we only ever push new descriptions to the back and pop from the front.
1210            let mut minted_batches = VecDeque::new();
1211
1212            // Once we start seeing new data, we'll roll everything into a single catch-up commit
1213            // before beginning the steady state commit interval.
1214            let catchup_start = if *resume_upper == [Timestamp::minimum()] {
1215                // If we're hydrating from a source snapshot, we immediately emit the snapshot's batch description.
1216                let batch_upper = Antichain::from_elem(
1217                    as_of.as_option().expect("as_of not empty").step_forward());
1218                let batch = (as_of.clone(), batch_upper.clone());
1219                minted_batches.push_back(batch.clone());
1220                batch_desc_output.give(&capset[0], batch);
1221                capset.downgrade(batch_upper.clone());
1222
1223                // The "catch-up" batch starts at the end of the snapshot batch.
1224                batch_upper
1225            } else {
1226                // If we're resuming, the "catch-up" batch starts at the start of data, i.e. resume_upper.
1227                resume_upper.clone()
1228            };
1229
1230            loop {
1231                if let Some(event) = input.next().await {
1232                    match event {
1233                        Event::Data(_, _) => continue,
1234                        Event::Progress(frontier) => {
1235                            observed_frontier = frontier;
1236                        }
1237                    }
1238                } else {
1239                    return Ok(());
1240                }
1241
1242                if !initialized {
1243                    if observed_frontier.is_empty() {
1244                        // Bounded inputs can close (frontier becomes empty) before we finish
1245                        // initialization. For example, a loadgen source configured for a finite
1246                        // dataset may emit all rows at time t and then immediately close.
1247                        // Mint one final batch with an empty upper. The input is closed, so
1248                        // that batch covers all remaining data on every worker, and committing
1249                        // it records the sink as complete.
1250                        if catchup_start.is_empty() {
1251                            // A previous incarnation already committed through the empty
1252                            // frontier. Nothing left to do.
1253                            return Ok(());
1254                        }
1255                        debug!(
1256                            ?sink_id,
1257                            %name_for_logging,
1258                            batch_lower = %catchup_start.pretty(),
1259                            "iceberg mint input closed before initialization; minting final batch"
1260                        );
1261                        let batch = (catchup_start.clone(), Antichain::new());
1262                        batch_desc_output.give(&capset[0], batch);
1263                        return Ok(());
1264                    }
1265
1266                    // Don't make empty commits while we wait ^for the first data to be ready.
1267                    // (^for the frontier to indicate there _could_ be data ready)
1268                    if !PartialOrder::less_than(&catchup_start, &observed_frontier)
1269                    {
1270                        continue;
1271                    }
1272
1273                    let mut batch_descriptions = vec![];
1274                    let mut current_upper = observed_frontier.clone();
1275                    let current_upper_ts = observed_frontier.as_option().expect("frontier not empty").clone();
1276                    debug!(
1277                        ?sink_id,
1278                        %name_for_logging,
1279                        batch_lower = %catchup_start.pretty(),
1280                        current_upper = %current_upper.pretty(),
1281                        "iceberg mint initializing (catch-up batch)"
1282                    );
1283                    debug!(
1284                        "{}: creating catch-up batch [{}, {})",
1285                        name_for_logging,
1286                        catchup_start.pretty(),
1287                        current_upper.pretty()
1288                    );
1289                    batch_descriptions.push((catchup_start.clone(), current_upper.clone()));
1290
1291                    // Mint initial future batch descriptions at configurable intervals
1292                    for i in 1..INITIAL_DESCRIPTIONS_TO_MINT {
1293                        let duration_millis = commit_interval.as_millis()
1294                            .checked_mul(u128::from(i))
1295                            .expect("commit interval multiplication overflow");
1296                        let duration_ts = Timestamp::new(
1297                            u64::try_from(duration_millis)
1298                                .expect("commit interval too large for u64"),
1299                        );
1300                        let desired_batch_upper = Antichain::from_elem(
1301                            current_upper_ts.step_forward_by(&duration_ts),
1302                        );
1303
1304                        let batch_description =
1305                            (current_upper.clone(), desired_batch_upper.clone());
1306                        debug!(
1307                            "{}: minting future batch {}/{} [{}, {})",
1308                            name_for_logging,
1309                            i,
1310                            INITIAL_DESCRIPTIONS_TO_MINT,
1311                            current_upper.pretty(),
1312                            desired_batch_upper.pretty()
1313                        );
1314                        current_upper = batch_description.1.clone();
1315                        batch_descriptions.push(batch_description);
1316                    }
1317
1318                    minted_batches.extend(batch_descriptions.clone());
1319
1320                    for desc in batch_descriptions {
1321                        batch_desc_output.give(&capset[0], desc);
1322                    }
1323
1324                    capset.downgrade(current_upper);
1325
1326                    initialized = true;
1327                } else {
1328                    if observed_frontier.is_empty() {
1329                        // We're done!
1330                        return Ok(());
1331                    }
1332                    // Maintain a sliding window: when the oldest batch becomes ready, retire it
1333                    // and mint a new future batch to keep the pipeline full
1334                    while let Some(oldest_desc) = minted_batches.front() {
1335                        let oldest_upper = &oldest_desc.1;
1336                        if !PartialOrder::less_equal(oldest_upper, &observed_frontier) {
1337                            break;
1338                        }
1339
1340                        let newest_upper = minted_batches.back().unwrap().1.clone();
1341                        let new_lower = newest_upper.clone();
1342                        let duration_ts = Timestamp::new(commit_interval.as_millis()
1343                            .try_into()
1344                            .expect("commit interval too large for u64"));
1345                        let new_upper = Antichain::from_elem(newest_upper
1346                            .as_option()
1347                            .unwrap()
1348                            .step_forward_by(&duration_ts));
1349
1350                        let new_batch_description = (new_lower.clone(), new_upper.clone());
1351                        minted_batches.pop_front();
1352                        minted_batches.push_back(new_batch_description.clone());
1353
1354                        batch_desc_output.give(&capset[0], new_batch_description);
1355
1356                        capset.downgrade(new_upper);
1357                    }
1358                }
1359            }
1360        })
1361    });
1362
1363    let statuses = errors.map(|error| HealthStatusMessage {
1364        id: None,
1365        update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
1366        namespace: StatusNamespace::Iceberg,
1367    });
1368    (
1369        batch_desc_stream,
1370        table_ready_stream,
1371        statuses,
1372        button.press_on_drop(),
1373    )
1374}
1375
1376#[derive(Clone, Debug, Serialize, Deserialize)]
1377#[serde(try_from = "AvroDataFile", into = "AvroDataFile")]
1378struct SerializableDataFile {
1379    pub data_file: DataFile,
1380    pub schema: Schema,
1381}
1382
1383/// A wrapper around Iceberg's DataFile that implements Serialize and Deserialize.
1384/// This is slightly complicated by the fact that Iceberg's DataFile doesn't implement
1385/// these traits directly, so we serialize to/from Avro bytes (which Iceberg supports natively).
1386/// The avro ser(de) also requires the Iceberg schema to be provided, so we include that as well.
1387/// It is distinctly possible that this is overkill, but it avoids re-implementing
1388/// Iceberg's serialization logic here.
1389/// If at some point this becomes a serious overhead, we can revisit this decision.
1390#[derive(Clone, Debug, Serialize, Deserialize)]
1391struct AvroDataFile {
1392    pub data_file: Vec<u8>,
1393    /// Schema serialized as JSON bytes to avoid bincode issues with HashMap
1394    pub schema: Vec<u8>,
1395}
1396
1397impl From<SerializableDataFile> for AvroDataFile {
1398    fn from(value: SerializableDataFile) -> Self {
1399        let mut data_file = Vec::new();
1400        write_data_files_to_avro(
1401            &mut data_file,
1402            [value.data_file],
1403            &StructType::new(vec![]),
1404            FormatVersion::V2,
1405        )
1406        .expect("serialization into buffer");
1407        let schema = serde_json::to_vec(&value.schema).expect("schema serialization");
1408        AvroDataFile { data_file, schema }
1409    }
1410}
1411
1412impl TryFrom<AvroDataFile> for SerializableDataFile {
1413    type Error = String;
1414
1415    fn try_from(value: AvroDataFile) -> Result<Self, Self::Error> {
1416        let schema: Schema = serde_json::from_slice(&value.schema)
1417            .map_err(|e| format!("Failed to deserialize schema: {}", e))?;
1418        let data_files = read_data_files_from_avro(
1419            &mut &*value.data_file,
1420            &schema,
1421            0,
1422            &StructType::new(vec![]),
1423            FormatVersion::V2,
1424        )
1425        .map_err_to_string_with_causes()?;
1426        let Some(data_file) = data_files.into_iter().next() else {
1427            return Err("No DataFile found in Avro data".into());
1428        };
1429        Ok(SerializableDataFile { data_file, schema })
1430    }
1431}
1432
1433/// A DataFile along with its associated batch description (lower and upper bounds).
1434#[derive(Clone, Debug, Serialize, Deserialize)]
1435struct BoundedDataFile {
1436    pub data_file: SerializableDataFile,
1437    pub batch_desc: (Antichain<Timestamp>, Antichain<Timestamp>),
1438}
1439
1440impl BoundedDataFile {
1441    pub fn new(
1442        file: DataFile,
1443        schema: Schema,
1444        batch_desc: (Antichain<Timestamp>, Antichain<Timestamp>),
1445    ) -> Self {
1446        Self {
1447            data_file: SerializableDataFile {
1448                data_file: file,
1449                schema,
1450            },
1451            batch_desc,
1452        }
1453    }
1454
1455    pub fn batch_desc(&self) -> &(Antichain<Timestamp>, Antichain<Timestamp>) {
1456        &self.batch_desc
1457    }
1458
1459    pub fn data_file(&self) -> &DataFile {
1460        &self.data_file.data_file
1461    }
1462
1463    pub fn into_data_file(self) -> DataFile {
1464        self.data_file.data_file
1465    }
1466}
1467
1468/// A set of DataFiles along with their associated batch descriptions.
1469#[derive(Clone, Debug, Default)]
1470struct BoundedDataFileSet {
1471    pub data_files: Vec<BoundedDataFile>,
1472}
1473
1474/// Construct the envelope-specific closures that [`write_data_files`] needs.
1475///
1476/// Write rows into Parquet data files bounded by batch descriptions.
1477/// Rows are matched to batches by timestamp; if a batch description hasn't arrived yet,
1478/// rows are stashed until it does. This allows batches to be minted ahead of data arrival.
1479fn write_data_files<'scope, H: EnvelopeHandler + 'static>(
1480    name: String,
1481    input: SinkBatchStream<'scope>,
1482    batch_desc_input: StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
1483    table_ready_stream: StreamVec<'scope, Timestamp, Infallible>,
1484    sink_id: GlobalId,
1485    from_id: GlobalId,
1486    key_is_synthetic: bool,
1487    as_of: Antichain<Timestamp>,
1488    connection: IcebergSinkConnection,
1489    storage_configuration: StorageConfiguration,
1490    materialize_arrow_schema: Arc<ArrowSchema>,
1491    metrics: Arc<IcebergSinkMetrics>,
1492    statistics: SinkStatistics,
1493) -> (
1494    StreamVec<'scope, Timestamp, BoundedDataFile>,
1495    StreamVec<'scope, Timestamp, HealthStatusMessage>,
1496    PressOnDropButton,
1497) {
1498    let scope = input.scope();
1499    let name_for_logging = name.clone();
1500    let mut builder = OperatorBuilder::new(name, scope.clone());
1501
1502    let (output, output_stream) = builder.new_output::<CapacityContainerBuilder<_>>();
1503
1504    let mut table_ready_input = builder.new_disconnected_input(table_ready_stream, Pipeline);
1505    let mut batch_desc_input =
1506        builder.new_input_for(batch_desc_input.broadcast(), Pipeline, &output);
1507    let mut input = builder.new_disconnected_input(input, Pipeline);
1508
1509    let (button, errors): (_, StreamVec<'scope, Timestamp, Rc<anyhow::Error>>) = builder
1510        .build_fallible(move |caps| {
1511            Box::pin(async move {
1512                let [capset]: &mut [_; 1] = caps.try_into().unwrap();
1513                let catalog = connection
1514                    .catalog_connection
1515                    .connect(&storage_configuration, InTask::Yes)
1516                    .await
1517                    .with_context(|| {
1518                        format!(
1519                            "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
1520                            connection.catalog_connection.uri,
1521                            connection.namespace,
1522                            connection.table
1523                        )
1524                    })?;
1525
1526                let namespace_ident = NamespaceIdent::new(connection.namespace.clone());
1527                let table_ident = TableIdent::new(namespace_ident, connection.table.clone());
1528                while let Some(_) = table_ready_input.next().await {
1529                    // Wait for table to be ready
1530                }
1531                let table = catalog.load_table(&table_ident).await.with_context(|| {
1532                    format!(
1533                        "Failed to load Iceberg table '{}.{}' in write_data_files operator",
1534                        connection.namespace, connection.table
1535                    )
1536                })?;
1537
1538                let table_metadata = table.metadata().clone();
1539                let current_schema = Arc::clone(table_metadata.current_schema());
1540
1541                // Merge Materialize extension metadata into the Iceberg schema.
1542                // We need extension metadata for ArrowBuilder to work correctly (it uses
1543                // extension names to know how to handle different types like records vs arrays).
1544                let arrow_schema = Arc::new(
1545                    merge_materialize_metadata_into_iceberg_schema(
1546                        materialize_arrow_schema.as_ref(),
1547                        current_schema.as_ref(),
1548                    )
1549                    .context("Failed to merge Materialize metadata into Iceberg schema")?,
1550                );
1551
1552                // WORKAROUND: S3 Tables catalog incorrectly sets location to the metadata file path
1553                // instead of the warehouse root. Strip off the /metadata/*.metadata.json suffix.
1554                // No clear way to detect this properly right now, so we use heuristics.
1555                let location = table_metadata.location();
1556                let corrected_location = match location.rsplit_once("/metadata/") {
1557                    Some((a, b)) if b.ends_with(".metadata.json") => a,
1558                    _ => location,
1559                };
1560
1561                let data_location = format!("{}/data", corrected_location);
1562                let location_generator =
1563                    DefaultLocationGenerator::with_data_location(data_location);
1564
1565                // Add a unique suffix to avoid filename collisions across restarts and workers
1566                let unique_suffix = format!("-{}", uuid::Uuid::new_v4());
1567                let file_name_generator = DefaultFileNameGenerator::new(
1568                    PARQUET_FILE_PREFIX.to_string(),
1569                    Some(unique_suffix),
1570                    iceberg::spec::DataFileFormat::Parquet,
1571                );
1572
1573                let file_io = table.file_io().clone();
1574
1575                let writer_properties = WriterProperties::new();
1576
1577                let ctx = WriterContext {
1578                    arrow_schema,
1579                    current_schema: Arc::clone(&current_schema),
1580                    file_io,
1581                    location_generator,
1582                    file_name_generator,
1583                    writer_properties,
1584                };
1585                let handler = H::new(ctx, &connection, &materialize_arrow_schema)?;
1586                let mut pk_warner =
1587                    (!key_is_synthetic).then(|| PkViolationWarner::new(sink_id, from_id));
1588
1589                // Rows can arrive before their batch description due to dataflow parallelism.
1590                // Stash them until we know which batch they belong to.
1591                // Keyed by the lower bound (per arrangement batch) of the rows.
1592                let mut stashed_rows: VecDeque<ArcBatch<OrdValBatch<_>>> = VecDeque::new();
1593
1594                // Track batches currently being written. When a row arrives, we check if it belongs
1595                // to an in-flight batch. When frontiers advance to a batch's upper, we close the
1596                // writer and emit its data files downstream.
1597                let mut in_flight_batches: VecDeque<(
1598                    (Antichain<Timestamp>, Antichain<Timestamp>),
1599                    Box<dyn IcebergWriter>,
1600                )> = VecDeque::new();
1601
1602                // The bounds of the most recently received batch description and input batch.
1603                // `with_ready_batches` relies on both inputs arriving in order and
1604                // non-overlapping. These track that invariant for the checks below.
1605                let mut last_batch_desc: Option<BatchDescription> = None;
1606                let mut last_input_bounds: Option<(Antichain<Timestamp>, Antichain<Timestamp>)> =
1607                    None;
1608
1609                let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
1610                let mut input_frontier = Antichain::from_elem(Timestamp::minimum());
1611
1612                while !(batch_description_frontier.is_empty() && input_frontier.is_empty()) {
1613                    tokio::select! {
1614                        _ = batch_desc_input.ready() => {},
1615                        _ = input.ready() => {}
1616                    }
1617
1618                    // Operator Recipe Step 1: Read all the input.
1619
1620                    // Read all the incoming batch descriptions.
1621                    while let Some(event) = batch_desc_input.next_sync() {
1622                        match event {
1623                            Event::Data(_cap, data) => {
1624                                for batch_desc in data {
1625                                    let (lower, upper) = &batch_desc;
1626
1627                                    if let Some((prev_lower, prev_upper)) = last_batch_desc.as_ref()
1628                                    {
1629                                        if prev_upper != lower {
1630                                            anyhow::bail!(
1631                                                "batch descriptions must arrive in order, non-overlapping, \
1632                                                and without gaps: previous [{}, {}), new [{}, {})",
1633                                                prev_lower.pretty(),
1634                                                prev_upper.pretty(),
1635                                                lower.pretty(),
1636                                                upper.pretty(),
1637                                            );
1638                                        }
1639                                    }
1640                                    last_batch_desc = Some(batch_desc.clone());
1641
1642                                    // Disable seen_rows tracking for snapshot batch to save memory
1643                                    let is_snapshot = lower == &as_of;
1644                                    debug!(
1645                                        "{}: received batch description [{}, {}), snapshot={}",
1646                                        name_for_logging,
1647                                        lower.pretty(),
1648                                        upper.pretty(),
1649                                        is_snapshot
1650                                    );
1651                                    let batch_writer = handler.create_writer(is_snapshot).await?;
1652                                    in_flight_batches.push_back((batch_desc.clone(), batch_writer));
1653                                }
1654                            }
1655                            Event::Progress(frontier) => {
1656                                batch_description_frontier = frontier;
1657                            }
1658                        }
1659                    }
1660
1661                    // Read all the incoming (arrangement batches of) rows.
1662                    while let Some(event) = input.next_sync() {
1663                        match event {
1664                            Event::Data(_cap, data) => {
1665                                for rows in &data {
1666                                    if let Some((prev_lower, prev_upper)) =
1667                                        last_input_bounds.as_ref()
1668                                    {
1669                                        // If the collection doesn't change and the frontier advances,
1670                                        // we can (correctly) observe gaps between input batches.
1671                                        // This differs from output batch descriptions, which are constructed without gaps.
1672                                        if !PartialOrder::less_equal(prev_upper, rows.lower()) {
1673                                            anyhow::bail!(
1674                                                "input batches must arrive in order and \
1675                                                non-overlapping: previous [{}, {}), new [{}, {})",
1676                                                prev_lower.pretty(),
1677                                                prev_upper.pretty(),
1678                                                rows.lower().pretty(),
1679                                                rows.upper().pretty(),
1680                                            );
1681                                        }
1682                                    }
1683                                    last_input_bounds =
1684                                        Some((rows.lower().clone(), rows.upper().clone()));
1685
1686                                    stashed_rows.push_back(rows.clone());
1687                                }
1688                            }
1689                            Event::Progress(frontier) => {
1690                                input_frontier = frontier;
1691                            }
1692                        }
1693                    }
1694
1695                    metrics.stashed_rows.set(u64::cast_from(
1696                        stashed_rows.iter().map(|rows| rows.len()).sum::<usize>(),
1697                    ));
1698
1699                    // Operator Recipe Steps 2-4: Consult frontiers. Plan work. Do all the work.
1700
1701                    // Report staged messages periodically during writes so progress is
1702                    // visible while a large batch is still open.
1703                    let mut staged_messages_since_flush: u64 = 0;
1704
1705                    // How to write rows from a(n arrangement) batch into a(n Iceberg) batch.
1706                    let write_rows = async |rows: &OrdValBatch<_>,
1707                                            (lower, upper): BatchDescription,
1708                                            batch_writer: &mut Box<dyn IcebergWriter>|
1709                           -> Result<(), anyhow::Error> {
1710                        for_each_diff_pair_async(
1711                            rows,
1712                            Some(lower),
1713                            Some(upper),
1714                            async |key, time, diff_pair| -> Result<(), anyhow::Error> {
1715                                if let Some(warner) = pk_warner.as_mut() {
1716                                    warner.observe(key, time);
1717                                }
1718
1719                                let record_batch = handler
1720                                    .row_to_batch(diff_pair, time)
1721                                    .context("failed to convert row to recordbatch")?;
1722                                staged_messages_since_flush +=
1723                                    u64::cast_from(record_batch.num_rows());
1724                                batch_writer
1725                                    .write(record_batch)
1726                                    .await
1727                                    .context("failed to write recordbatch")?;
1728                                if staged_messages_since_flush >= 10_000 {
1729                                    statistics.inc_messages_staged_by(staged_messages_since_flush);
1730                                    staged_messages_since_flush = 0;
1731                                }
1732                                Ok(())
1733                            },
1734                        )
1735                        .await?;
1736                        // Flush after each batch so the final `(key, time)` group of the walk is
1737                        // resolved immediately — a PK violation in the last group is otherwise held
1738                        // until more data arrives or the operator shuts down.
1739                        if let Some(warner) = pk_warner.as_mut() {
1740                            warner.flush();
1741                        }
1742                        Ok(())
1743                    };
1744
1745                    // How to seal the data files for an Iceberg commit.
1746                    let close_batch = async |batch_desc: BatchDescription,
1747                                             batch_writer: &mut Box<dyn IcebergWriter>|
1748                           -> Result<(), anyhow::Error> {
1749                        let close_started_at = Instant::now();
1750                        let data_files = batch_writer.close().await;
1751                        metrics
1752                            .writer_close_duration_seconds
1753                            .observe(close_started_at.elapsed().as_secs_f64());
1754                        let data_files = data_files.context("Failed to close batch writer")?;
1755                        debug!(
1756                            "{}: closed batch [{}, {}), wrote {} files",
1757                            name_for_logging,
1758                            batch_desc.0.pretty(),
1759                            batch_desc.1.pretty(),
1760                            data_files.len()
1761                        );
1762                        for data_file in data_files {
1763                            match data_file.content_type() {
1764                                iceberg::spec::DataContentType::Data => {
1765                                    metrics.data_files_written.inc();
1766                                }
1767                                iceberg::spec::DataContentType::PositionDeletes
1768                                | iceberg::spec::DataContentType::EqualityDeletes => {
1769                                    metrics.delete_files_written.inc();
1770                                }
1771                            }
1772                            statistics.inc_bytes_staged_by(data_file.file_size_in_bytes());
1773                            let file = BoundedDataFile::new(
1774                                data_file,
1775                                current_schema.as_ref().clone(),
1776                                batch_desc.clone(),
1777                            );
1778                            output.give(&capset[0], file);
1779                        }
1780
1781                        // Operator Recipe Step 5: Downgrade or drop capabilities.
1782
1783                        capset.downgrade(batch_desc.1.clone());
1784                        Ok(())
1785                    };
1786
1787                    // Write the rows and seal the data files.
1788                    with_ready_batches(
1789                        input_frontier.clone(),
1790                        &mut stashed_rows,
1791                        batch_description_frontier.clone(),
1792                        &mut in_flight_batches,
1793                        write_rows,
1794                        close_batch,
1795                    )
1796                    .await?;
1797
1798                    if staged_messages_since_flush > 0 {
1799                        statistics.inc_messages_staged_by(staged_messages_since_flush);
1800                    }
1801                    metrics.stashed_rows.set(u64::cast_from(
1802                        stashed_rows.iter().map(|rows| rows.len()).sum::<usize>(),
1803                    ));
1804                }
1805                Ok(())
1806            })
1807        });
1808
1809    let statuses = errors.map(|error| HealthStatusMessage {
1810        id: None,
1811        update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
1812        namespace: StatusNamespace::Iceberg,
1813    });
1814    (output_stream, statuses, button.press_on_drop())
1815}
1816
1817/// The `[lower, upper)` frontier bounds of one Iceberg commit.
1818type BatchDescription = (Antichain<Timestamp>, Antichain<Timestamp>);
1819
1820/// Write out as much of the input as we can.
1821///
1822/// Drop input batches when:
1823/// - their contents have all been written out
1824/// - no possible future output batch could need their contents
1825///
1826/// Close and drop output batches when:
1827/// - no possible future input batch could overlap with their time window
1828///
1829/// Invariant: We assume the batches in each stream (input vs output)
1830/// are in order and non-overlapping.
1831async fn with_ready_batches<L: Layout, W, Write, Close>(
1832    input_frontier: Antichain<Timestamp>,
1833    input_batches: &mut VecDeque<ArcBatch<OrdValBatch<L>>>,
1834    output_frontier: Antichain<Timestamp>,
1835    output_batches: &mut VecDeque<(BatchDescription, W)>,
1836    mut write_rows: Write,
1837    mut close_batch: Close,
1838) -> Result<(), anyhow::Error>
1839where
1840    L::TimeContainer: BatchContainer<Owned = Timestamp>,
1841    Write: AsyncFnMut(&OrdValBatch<L>, BatchDescription, &mut W) -> Result<(), anyhow::Error>,
1842    Close: AsyncFnMut(BatchDescription, &mut W) -> Result<(), anyhow::Error>,
1843{
1844    loop {
1845        {
1846            // Drop any input batches that fall below the lowest output batch.
1847            // No future output batch could need these inputs.
1848            let output_lower = output_batches
1849                .front()
1850                .map_or(&output_frontier, |((lower, _), _)| lower);
1851            while input_batches
1852                .pop_front_if(|rows| PartialOrder::less_equal(rows.upper(), output_lower))
1853                .is_some()
1854            {}
1855        }
1856
1857        {
1858            // Close and drop any output batches that fall below the lowest input batch.
1859            // No future inputs can arrive for these batches.
1860            let input_lower = input_batches
1861                .front()
1862                .map_or(&input_frontier, |rows| rows.lower());
1863            while let Some((batch_desc, mut batch_writer)) =
1864                output_batches.pop_front_if(|((_, batch_upper), _)| {
1865                    PartialOrder::less_equal(batch_upper, input_lower)
1866                })
1867            {
1868                close_batch(batch_desc, &mut batch_writer).await?;
1869            }
1870        }
1871
1872        let Some((batch_desc, batch_writer)) = output_batches.front_mut() else {
1873            // We're still waiting for descriptions of batches to write to.
1874            break;
1875        };
1876
1877        let Some(rows) = input_batches.front() else {
1878            // We're still waiting for rows to write.
1879            break;
1880        };
1881
1882        // If there were no overlap between the lowest input batch and the lowest output batch,
1883        // we'd have dropped the lower one already.
1884        // Since we still have both a lowest input batch and a lowest output batch,
1885        // there must be overlap.
1886
1887        // Write (the relevant portion of) the lowest input batch to the lowest output batch.
1888        // Drop whichever one's "upper" comes first. If they end simultaneously, drop both.
1889        write_rows(rows, batch_desc.clone(), batch_writer).await?;
1890        let output_upper = batch_desc.1.clone();
1891        let rows_upper = rows.upper();
1892        if PartialOrder::less_equal(&output_upper, rows_upper) {
1893            // Close and drop the output batch.
1894            let (batch_desc, mut batch_writer) =
1895                output_batches.pop_front().expect("already checked front");
1896            close_batch(batch_desc, &mut batch_writer).await?;
1897        }
1898        if PartialOrder::less_equal(rows_upper, &output_upper) {
1899            // Drop the input batch.
1900            input_batches.pop_front();
1901        }
1902
1903        // At least one of the two conditions above must be true,
1904        // so every loop iteration shrinks working set (of input/output batches).
1905        // Therefore, this loop must terminate.
1906    }
1907
1908    Ok(())
1909}
1910
1911#[cfg(test)]
1912mod tests {
1913    use super::*;
1914    use iceberg::spec::{PrimitiveType, Type};
1915    use mz_repr::SqlScalarType;
1916    use mz_storage_types::sinks::ICEBERG_UINT64_DECIMAL_PRECISION;
1917
1918    #[mz_ore::test]
1919    fn test_iceberg_type_overrides() {
1920        // UInt16 should override to Int32
1921        let result = iceberg_type_overrides(&SqlScalarType::UInt16);
1922        assert_eq!(result.unwrap().0, DataType::Int32);
1923
1924        // UInt32 should override to Int64
1925        let result = iceberg_type_overrides(&SqlScalarType::UInt32);
1926        assert_eq!(result.unwrap().0, DataType::Int64);
1927
1928        // UInt64 should override to Decimal128(20, 0)
1929        let result = iceberg_type_overrides(&SqlScalarType::UInt64);
1930        assert_eq!(
1931            result.unwrap().0,
1932            DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
1933        );
1934
1935        // MzTimestamp should override to Decimal128(20, 0)
1936        let result = iceberg_type_overrides(&SqlScalarType::MzTimestamp);
1937        assert_eq!(
1938            result.unwrap().0,
1939            DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
1940        );
1941
1942        // Other types should return None (use default)
1943        assert!(iceberg_type_overrides(&SqlScalarType::Int32).is_none());
1944        assert!(iceberg_type_overrides(&SqlScalarType::String).is_none());
1945        assert!(iceberg_type_overrides(&SqlScalarType::Bool).is_none());
1946    }
1947
1948    #[mz_ore::test]
1949    fn test_iceberg_schema_with_nested_uint64() {
1950        // Test that desc_to_schema_with_overrides handles nested UInt64
1951        // by using iceberg_type_overrides which applies recursively
1952        let desc = mz_repr::RelationDesc::builder()
1953            .with_column(
1954                "items",
1955                SqlScalarType::List {
1956                    element_type: Box::new(SqlScalarType::UInt64),
1957                    custom_id: None,
1958                }
1959                .nullable(true),
1960            )
1961            .finish();
1962
1963        let schema =
1964            mz_arrow_util::builder::desc_to_schema_with_overrides(&desc, iceberg_type_overrides)
1965                .expect("schema conversion should succeed");
1966
1967        // The inner element should be Decimal128, not UInt64
1968        if let DataType::List(field) = schema.field(0).data_type() {
1969            assert_eq!(
1970                field.data_type(),
1971                &DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
1972            );
1973        } else {
1974            panic!("Expected List type");
1975        }
1976    }
1977
1978    #[mz_ore::test]
1979    fn test_iceberg_interval_override() {
1980        // Interval should override to LargeUtf8 (string) for Iceberg
1981        let result = iceberg_type_overrides(&SqlScalarType::Interval);
1982        assert_eq!(result.unwrap().0, DataType::LargeUtf8);
1983
1984        // Test full schema conversion with interval column
1985        let desc = mz_repr::RelationDesc::builder()
1986            .with_column("id", SqlScalarType::Int32.nullable(false))
1987            .with_column("dur", SqlScalarType::Interval.nullable(true))
1988            .finish();
1989
1990        let (arrow_schema, iceberg_schema) =
1991            relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
1992
1993        // Arrow schema should have LargeUtf8 for interval
1994        assert_eq!(arrow_schema.field(1).data_type(), &DataType::LargeUtf8);
1995
1996        // Iceberg schema should have String type
1997        let field = iceberg_schema
1998            .as_struct()
1999            .field_by_name("dur")
2000            .expect("field should exist");
2001        assert_eq!(*field.field_type, Type::Primitive(PrimitiveType::String));
2002    }
2003
2004    #[mz_ore::test]
2005    fn test_iceberg_range_schema() {
2006        // Test full schema conversion with range column
2007        let desc = mz_repr::RelationDesc::builder()
2008            .with_column("id", SqlScalarType::Int32.nullable(false))
2009            .with_column(
2010                "r",
2011                SqlScalarType::Range {
2012                    element_type: Box::new(SqlScalarType::Int32),
2013                }
2014                .nullable(true),
2015            )
2016            .finish();
2017
2018        let (_arrow_schema, iceberg_schema) =
2019            relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2020
2021        // Iceberg schema should have a struct type for the range
2022        let field = iceberg_schema
2023            .as_struct()
2024            .field_by_name("r")
2025            .expect("field should exist");
2026        assert!(
2027            matches!(&*field.field_type, Type::Struct(_)),
2028            "range should be struct, got: {:?}",
2029            field.field_type
2030        );
2031    }
2032
2033    #[mz_ore::test]
2034    fn equality_ids_follow_iceberg_field_ids() {
2035        let map_entries = Field::new(
2036            "entries",
2037            DataType::Struct(
2038                vec![
2039                    Field::new("key", DataType::Utf8, false),
2040                    Field::new("value", DataType::Utf8, true),
2041                ]
2042                .into(),
2043            ),
2044            false,
2045        );
2046        let materialize_arrow_schema = ArrowSchema::new(vec![
2047            Field::new("attrs", DataType::Map(Arc::new(map_entries), false), true),
2048            Field::new("key_col", DataType::Int32, false),
2049        ]);
2050        let materialize_arrow_schema = add_field_ids_to_arrow_schema(materialize_arrow_schema);
2051        let iceberg_schema = arrow_schema_to_schema(&materialize_arrow_schema)
2052            .expect("schema conversion should succeed");
2053
2054        let equality_ids =
2055            equality_ids_for_indices(&iceberg_schema, &materialize_arrow_schema, &[1])
2056                .expect("field lookup should succeed");
2057
2058        let expected_id = iceberg_schema
2059            .as_struct()
2060            .field_by_name("key_col")
2061            .expect("top-level field should exist")
2062            .id;
2063        assert_eq!(equality_ids, vec![expected_id]);
2064        assert_ne!(expected_id, 2);
2065    }
2066
2067    /// Regression test: iceberg-rust names map fields `key_value`/`key`/`value`
2068    /// while Materialize uses `entries`/`keys`/`values`. The schema merge must
2069    /// still copy the value field's extension metadata across so ArrowBuilder
2070    /// can build the inner builder.
2071    #[mz_ore::test]
2072    #[allow(clippy::disallowed_types)]
2073    fn merge_map_entries_preserves_value_extension_metadata() {
2074        use std::collections::HashMap;
2075
2076        let mz_value_metadata = HashMap::from([(
2077            ARROW_EXTENSION_NAME_KEY.to_string(),
2078            "materialize.v1.string".to_string(),
2079        )]);
2080        let mz_entries = Field::new(
2081            "entries",
2082            DataType::Struct(
2083                vec![
2084                    Field::new("keys", DataType::Utf8, false),
2085                    Field::new("values", DataType::Utf8, true).with_metadata(mz_value_metadata),
2086                ]
2087                .into(),
2088            ),
2089            false,
2090        );
2091        let mz_map = Field::new("m", DataType::Map(Arc::new(mz_entries), false), true)
2092            .with_metadata(HashMap::from([(
2093                ARROW_EXTENSION_NAME_KEY.to_string(),
2094                "materialize.v1.map".to_string(),
2095            )]));
2096
2097        let iceberg_entries = Field::new(
2098            "key_value",
2099            DataType::Struct(
2100                vec![
2101                    Field::new("key", DataType::Utf8, false),
2102                    Field::new("value", DataType::Utf8, true),
2103                ]
2104                .into(),
2105            ),
2106            false,
2107        );
2108        let iceberg_map = Field::new("m", DataType::Map(Arc::new(iceberg_entries), false), true);
2109
2110        let merged = merge_field_metadata_recursive(&iceberg_map, Some(&mz_map))
2111            .expect("merge should succeed");
2112
2113        let entries = match merged.data_type() {
2114            DataType::Map(entries, _) => entries.as_ref(),
2115            other => panic!("expected Map, got {other:?}"),
2116        };
2117        let entry_fields = match entries.data_type() {
2118            DataType::Struct(fields) => fields,
2119            other => panic!("expected Struct, got {other:?}"),
2120        };
2121        // Iceberg naming must be preserved on the merged schema...
2122        assert_eq!(entry_fields[0].name(), "key");
2123        assert_eq!(entry_fields[1].name(), "value");
2124        // ...and the materialize extension must have been copied positionally
2125        // to the value field even though its name didn't match `values`.
2126        assert_eq!(
2127            entry_fields[1].metadata().get(ARROW_EXTENSION_NAME_KEY),
2128            Some(&"materialize.v1.string".to_string()),
2129        );
2130    }
2131
2132    mod with_ready_batches {
2133        use differential_dataflow::trace::Batch;
2134        use differential_dataflow::trace::implementations::Vector;
2135
2136        use super::*;
2137
2138        type TestBatch = OrdValBatch<Vector<((u64, u64), Timestamp, Diff)>>;
2139
2140        /// A frontier at `t`, or the empty (end-of-time) frontier for `None`.
2141        fn frontier(t: Option<u64>) -> Antichain<Timestamp> {
2142            t.map_or_else(Antichain::new, |t| Antichain::from_elem(Timestamp::new(t)))
2143        }
2144
2145        /// `[lower, upper)` bounds, with `None` for the empty upper.
2146        fn span(lower: u64, upper: Option<u64>) -> BatchDescription {
2147            (frontier(Some(lower)), frontier(upper))
2148        }
2149
2150        /// An input batch with the given bounds. The pairing logic under test
2151        /// only looks at bounds, so the batch holds no data.
2152        fn input(lower: u64, upper: Option<u64>) -> ArcBatch<TestBatch> {
2153            let (lower, upper) = span(lower, upper);
2154            ArcBatch(Arc::new(TestBatch::empty(lower, upper)))
2155        }
2156
2157        #[derive(Debug, PartialEq)]
2158        enum Call {
2159            /// (input batch bounds, output batch description)
2160            Write(BatchDescription, BatchDescription),
2161            Close(BatchDescription),
2162        }
2163
2164        /// Run `with_ready_batches` with recording callbacks and return the
2165        /// sequence of calls it made.
2166        async fn run(
2167            input_frontier: Antichain<Timestamp>,
2168            input_batches: &mut VecDeque<ArcBatch<TestBatch>>,
2169            output_frontier: Antichain<Timestamp>,
2170            output_batches: &mut VecDeque<(BatchDescription, ())>,
2171        ) -> Vec<Call> {
2172            let calls = RefCell::new(vec![]);
2173            with_ready_batches(
2174                input_frontier,
2175                input_batches,
2176                output_frontier,
2177                output_batches,
2178                async |rows: &TestBatch, desc, _writer: &mut ()| {
2179                    let bounds = (rows.lower().clone(), rows.upper().clone());
2180                    calls.borrow_mut().push(Call::Write(bounds, desc));
2181                    Ok(())
2182                },
2183                async |desc, _writer: &mut ()| {
2184                    calls.borrow_mut().push(Call::Close(desc));
2185                    Ok(())
2186                },
2187            )
2188            .await
2189            .expect("test callbacks never fail");
2190            calls.into_inner()
2191        }
2192
2193        #[mz_ore::test(tokio::test)]
2194        async fn input_batch_spanning_multiple_output_batches() {
2195            let mut inputs = VecDeque::from([input(0, Some(30))]);
2196            let mut outputs = VecDeque::from([
2197                (span(0, Some(10)), ()),
2198                (span(10, Some(20)), ()),
2199                (span(20, Some(30)), ()),
2200            ]);
2201
2202            let calls = run(
2203                frontier(Some(30)),
2204                &mut inputs,
2205                frontier(Some(30)),
2206                &mut outputs,
2207            )
2208            .await;
2209
2210            // The input batch is written once per overlapping output batch,
2211            // each of which closes as soon as the input covers its upper.
2212            assert_eq!(
2213                calls,
2214                vec![
2215                    Call::Write(span(0, Some(30)), span(0, Some(10))),
2216                    Call::Close(span(0, Some(10))),
2217                    Call::Write(span(0, Some(30)), span(10, Some(20))),
2218                    Call::Close(span(10, Some(20))),
2219                    Call::Write(span(0, Some(30)), span(20, Some(30))),
2220                    Call::Close(span(20, Some(30))),
2221                ]
2222            );
2223            assert!(inputs.is_empty());
2224            assert!(outputs.is_empty());
2225        }
2226
2227        #[mz_ore::test(tokio::test)]
2228        async fn output_batch_spanning_multiple_input_batches() {
2229            let mut inputs =
2230                VecDeque::from([input(0, Some(10)), input(10, Some(20)), input(20, Some(30))]);
2231            let mut outputs = VecDeque::from([(span(0, Some(30)), ())]);
2232
2233            let calls = run(
2234                frontier(Some(30)),
2235                &mut inputs,
2236                frontier(Some(30)),
2237                &mut outputs,
2238            )
2239            .await;
2240
2241            assert_eq!(
2242                calls,
2243                vec![
2244                    Call::Write(span(0, Some(10)), span(0, Some(30))),
2245                    Call::Write(span(10, Some(20)), span(0, Some(30))),
2246                    Call::Write(span(20, Some(30)), span(0, Some(30))),
2247                    Call::Close(span(0, Some(30))),
2248                ]
2249            );
2250            assert!(inputs.is_empty());
2251            assert!(outputs.is_empty());
2252        }
2253
2254        #[mz_ore::test(tokio::test)]
2255        async fn input_batch_retained_for_future_output_batches() {
2256            let mut inputs = VecDeque::from([input(0, Some(30))]);
2257            let mut outputs = VecDeque::from([(span(0, Some(10)), ())]);
2258
2259            let calls = run(
2260                frontier(Some(30)),
2261                &mut inputs,
2262                frontier(Some(10)),
2263                &mut outputs,
2264            )
2265            .await;
2266
2267            // The input batch extends past the only known output batch, so it
2268            // must stay queued for descriptions that haven't arrived yet.
2269            assert_eq!(
2270                calls,
2271                vec![
2272                    Call::Write(span(0, Some(30)), span(0, Some(10))),
2273                    Call::Close(span(0, Some(10))),
2274                ]
2275            );
2276            assert_eq!(inputs.len(), 1);
2277            assert!(outputs.is_empty());
2278        }
2279
2280        #[mz_ore::test(tokio::test)]
2281        async fn already_committed_input_batches_dropped_unwritten() {
2282            let mut inputs = VecDeque::from([input(0, Some(10)), input(10, Some(20))]);
2283            let mut outputs = VecDeque::from([(span(20, Some(30)), ())]);
2284
2285            let calls = run(
2286                frontier(Some(20)),
2287                &mut inputs,
2288                frontier(Some(30)),
2289                &mut outputs,
2290            )
2291            .await;
2292
2293            // Both input batches fall below the lowest output batch, so their
2294            // contents are already committed and they are dropped unwritten.
2295            // The output batch still waits for its own input.
2296            assert_eq!(calls, vec![]);
2297            assert!(inputs.is_empty());
2298            assert_eq!(outputs.len(), 1);
2299        }
2300
2301        #[mz_ore::test(tokio::test)]
2302        async fn output_batch_closes_empty_once_input_frontier_passes() {
2303            let mut outputs = VecDeque::from([(span(0, Some(10)), ())]);
2304
2305            // While the input frontier is short of the batch's upper, nothing
2306            // may close: rows for it could still arrive.
2307            let calls = run(
2308                frontier(Some(5)),
2309                &mut VecDeque::new(),
2310                frontier(Some(10)),
2311                &mut outputs,
2312            )
2313            .await;
2314            assert_eq!(calls, vec![]);
2315            assert_eq!(outputs.len(), 1);
2316
2317            // Once the input frontier reaches the upper, the batch closes
2318            // empty (an empty commit).
2319            let calls = run(
2320                frontier(Some(10)),
2321                &mut VecDeque::new(),
2322                frontier(Some(10)),
2323                &mut outputs,
2324            )
2325            .await;
2326            assert_eq!(calls, vec![Call::Close(span(0, Some(10)))]);
2327            assert!(outputs.is_empty());
2328        }
2329
2330        #[mz_ore::test(tokio::test)]
2331        async fn final_output_batch_with_empty_upper() {
2332            let mut inputs = VecDeque::from([input(20, Some(30))]);
2333            let mut outputs = VecDeque::from([(span(20, None), ())]);
2334
2335            // The sealing batch covers everything from 20 to the end of time.
2336            // It consumes all remaining input but only closes once the input
2337            // frontier is empty, i.e. the input is finished.
2338            let calls = run(
2339                frontier(Some(30)),
2340                &mut inputs,
2341                frontier(None),
2342                &mut outputs,
2343            )
2344            .await;
2345            assert_eq!(calls, vec![Call::Write(span(20, Some(30)), span(20, None))]);
2346            assert!(inputs.is_empty());
2347            assert_eq!(outputs.len(), 1);
2348
2349            let calls = run(frontier(None), &mut inputs, frontier(None), &mut outputs).await;
2350            assert_eq!(calls, vec![Call::Close(span(20, None))]);
2351            assert!(outputs.is_empty());
2352        }
2353    }
2354}
2355
2356/// Commit completed batches to Iceberg as snapshots.
2357/// Batches are committed in timestamp order to ensure strong consistency guarantees downstream.
2358/// Each snapshot includes the Materialize frontier in its metadata for resume support.
2359fn commit_to_iceberg<'scope>(
2360    name: String,
2361    sink_id: GlobalId,
2362    sink_version: u64,
2363    batch_input: StreamVec<'scope, Timestamp, BoundedDataFile>,
2364    batch_desc_input: StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
2365    table_ready_stream: StreamVec<'scope, Timestamp, Infallible>,
2366    write_frontier: Rc<RefCell<Antichain<Timestamp>>>,
2367    connection: IcebergSinkConnection,
2368    storage_configuration: StorageConfiguration,
2369    write_handle: impl Future<
2370        Output = anyhow::Result<WriteHandle<SourceData, (), Timestamp, StorageDiff>>,
2371    > + 'static,
2372    metrics: Arc<IcebergSinkMetrics>,
2373    statistics: SinkStatistics,
2374) -> (
2375    StreamVec<'scope, Timestamp, HealthStatusMessage>,
2376    PressOnDropButton,
2377) {
2378    let scope = batch_input.scope();
2379    let mut builder = OperatorBuilder::new(name, scope.clone());
2380
2381    let hashed_id = sink_id.hashed();
2382    let is_active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
2383    let name_for_logging = format!("{sink_id}-commit-to-iceberg");
2384
2385    let mut input = builder.new_disconnected_input(batch_input, Exchange::new(move |_| hashed_id));
2386    let mut batch_desc_input =
2387        builder.new_disconnected_input(batch_desc_input, Exchange::new(move |_| hashed_id));
2388    let mut table_ready_input = builder.new_disconnected_input(table_ready_stream, Pipeline);
2389
2390    let (button, errors) = builder.build_fallible(move |_caps| {
2391        Box::pin(async move {
2392            if !is_active_worker {
2393                write_frontier.borrow_mut().clear();
2394                return Ok(());
2395            }
2396
2397            let catalog = connection
2398                .catalog_connection
2399                .connect(&storage_configuration, InTask::Yes)
2400                .await
2401                .with_context(|| {
2402                    format!(
2403                        "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
2404                        connection.catalog_connection.uri, connection.namespace, connection.table
2405                    )
2406                })?;
2407
2408            let mut write_handle = write_handle.await?;
2409
2410            let namespace_ident = NamespaceIdent::new(connection.namespace.clone());
2411            let table_ident = TableIdent::new(namespace_ident, connection.table.clone());
2412            while let Some(_) = table_ready_input.next().await {
2413                // Wait for table to be ready
2414            }
2415            let mut table = catalog.load_table(&table_ident).await.with_context(|| {
2416                format!(
2417                    "Failed to load Iceberg table '{}.{}' in commit_to_iceberg operator",
2418                    connection.namespace, connection.table
2419                )
2420            })?;
2421
2422            #[allow(clippy::disallowed_types)]
2423            let mut batch_descriptions: std::collections::HashMap<
2424                (Antichain<Timestamp>, Antichain<Timestamp>),
2425                BoundedDataFileSet,
2426            > = std::collections::HashMap::new();
2427
2428            let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
2429            let mut input_frontier = Antichain::from_elem(Timestamp::minimum());
2430
2431            while !(batch_description_frontier.is_empty() && input_frontier.is_empty()) {
2432                tokio::select! {
2433                    _ = batch_desc_input.ready() => {},
2434                    _ = input.ready() => {}
2435                }
2436
2437                while let Some(event) = batch_desc_input.next_sync() {
2438                    match event {
2439                        Event::Data(_cap, data) => {
2440                            for batch_desc in data {
2441                                let prev = batch_descriptions
2442                                    .insert(batch_desc, BoundedDataFileSet { data_files: vec![] });
2443                                if let Some(prev) = prev {
2444                                    anyhow::bail!(
2445                                        "Duplicate batch description received \
2446                                         in commit operator: {:?}",
2447                                        prev
2448                                    );
2449                                }
2450                            }
2451                        }
2452                        Event::Progress(frontier) => {
2453                            batch_description_frontier = frontier;
2454                        }
2455                    }
2456                }
2457
2458                let ready_events = std::iter::from_fn(|| input.next_sync()).collect_vec();
2459                for event in ready_events {
2460                    match event {
2461                        Event::Data(_cap, data) => {
2462                            for bounded_data_file in data {
2463                                let entry = batch_descriptions
2464                                    .entry(bounded_data_file.batch_desc().clone())
2465                                    .or_default();
2466                                entry.data_files.push(bounded_data_file);
2467                            }
2468                        }
2469                        Event::Progress(frontier) => {
2470                            input_frontier = frontier;
2471                        }
2472                    }
2473                }
2474
2475                // Collect batches whose data files have all arrived.
2476                // The writer emits all data files for a batch at a capability <= the batch's
2477                // lower bound, then downgrades its capability to the batch's upper bound.
2478                // So once the input frontier advances past lower, we know the writer has
2479                // finished emitting files for this batch and dropped its capability.
2480                let mut done_batches: Vec<_> = batch_descriptions
2481                    .keys()
2482                    .filter(|(lower, _upper)| PartialOrder::less_than(lower, &input_frontier))
2483                    .cloned()
2484                    .collect();
2485
2486                // Commit batches in timestamp order to maintain consistency
2487                done_batches.sort_by(|a, b| {
2488                    if PartialOrder::less_than(a, b) {
2489                        Ordering::Less
2490                    } else if PartialOrder::less_than(b, a) {
2491                        Ordering::Greater
2492                    } else {
2493                        Ordering::Equal
2494                    }
2495                });
2496
2497                for batch in done_batches {
2498                    let file_set = batch_descriptions.remove(&batch).unwrap();
2499
2500                    let mut data_files = vec![];
2501                    let mut delete_files = vec![];
2502                    // Track totals for committed statistics
2503                    let mut total_messages: u64 = 0;
2504                    let mut total_bytes: u64 = 0;
2505                    for file in file_set.data_files {
2506                        total_messages += file.data_file().record_count();
2507                        total_bytes += file.data_file().file_size_in_bytes();
2508                        match file.data_file().content_type() {
2509                            iceberg::spec::DataContentType::Data => {
2510                                data_files.push(file.into_data_file());
2511                            }
2512                            iceberg::spec::DataContentType::PositionDeletes
2513                            | iceberg::spec::DataContentType::EqualityDeletes => {
2514                                delete_files.push(file.into_data_file());
2515                            }
2516                        }
2517                    }
2518
2519                    debug!(
2520                        ?sink_id,
2521                        %name_for_logging,
2522                        lower = %batch.0.pretty(),
2523                        upper = %batch.1.pretty(),
2524                        data_files = data_files.len(),
2525                        delete_files = delete_files.len(),
2526                        total_messages,
2527                        total_bytes,
2528                        "iceberg commit applying batch"
2529                    );
2530
2531                    let instant = Instant::now();
2532
2533                    let frontier = batch.1.clone();
2534                    let frontier_json = serde_json::to_string(&frontier.elements())
2535                        .context("Failed to serialize frontier to JSON")?;
2536                    let snapshot_properties = vec![
2537                        ("mz-sink-id".to_string(), sink_id.to_string()),
2538                        ("mz-frontier".to_string(), frontier_json),
2539                        ("mz-sink-version".to_string(), sink_version.to_string()),
2540                    ];
2541
2542                    let (table_state, commit_result) = Retry::default()
2543                        .max_tries(5)
2544                        .retry_async_with_state(table, |_, table| {
2545                            let snapshot_properties = snapshot_properties.clone();
2546                            let data_files = data_files.clone();
2547                            let delete_files = delete_files.clone();
2548                            let metrics = Arc::clone(&metrics);
2549                            let catalog = Arc::clone(&catalog);
2550                            let conn_namespace = connection.namespace.clone();
2551                            let conn_table = connection.table.clone();
2552                            let frontier = frontier.clone();
2553                            let batch_lower = batch.0.clone();
2554                            let batch_upper = batch.1.clone();
2555                            async move {
2556                                try_commit_batch(
2557                                    table,
2558                                    snapshot_properties,
2559                                    data_files,
2560                                    delete_files,
2561                                    catalog.as_ref(),
2562                                    &conn_namespace,
2563                                    &conn_table,
2564                                    sink_version,
2565                                    &frontier,
2566                                    &batch_lower,
2567                                    &batch_upper,
2568                                    &metrics,
2569                                )
2570                                .await
2571                            }
2572                        })
2573                        .await;
2574                    let commit_result = commit_result.with_context(|| {
2575                        format!(
2576                            "failed to commit batch to Iceberg table '{}.{}'",
2577                            connection.namespace, connection.table
2578                        )
2579                    });
2580                    table = table_state;
2581                    let duration = instant.elapsed();
2582                    metrics
2583                        .commit_duration_seconds
2584                        .observe(duration.as_secs_f64());
2585                    commit_result?;
2586
2587                    debug!(
2588                        ?sink_id,
2589                        %name_for_logging,
2590                        lower = %batch.0.pretty(),
2591                        upper = %batch.1.pretty(),
2592                        total_messages,
2593                        total_bytes,
2594                        ?duration,
2595                        "iceberg commit applied batch"
2596                    );
2597
2598                    metrics.snapshots_committed.inc();
2599                    statistics.inc_messages_committed_by(total_messages);
2600                    statistics.inc_bytes_committed_by(total_bytes);
2601
2602                    let mut expect_upper = write_handle.shared_upper();
2603                    loop {
2604                        if PartialOrder::less_equal(&frontier, &expect_upper) {
2605                            // The frontier has already been advanced as far as necessary.
2606                            break;
2607                        }
2608
2609                        const EMPTY: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
2610                        match write_handle
2611                            .compare_and_append(EMPTY, expect_upper, frontier.clone())
2612                            .await
2613                            .expect("valid usage")
2614                        {
2615                            Ok(()) => break,
2616                            Err(mismatch) => {
2617                                expect_upper = mismatch.current;
2618                            }
2619                        }
2620                    }
2621                    write_frontier.borrow_mut().clone_from(&frontier);
2622                }
2623            }
2624
2625            Ok(())
2626        })
2627    });
2628
2629    let statuses = errors.map(|error| HealthStatusMessage {
2630        id: None,
2631        update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
2632        namespace: StatusNamespace::Iceberg,
2633    });
2634
2635    (statuses, button.press_on_drop())
2636}
2637
2638impl<'scope> SinkRender<'scope> for IcebergSinkConnection {
2639    fn get_key_indices(&self) -> Option<&[usize]> {
2640        self.key_desc_and_indices
2641            .as_ref()
2642            .map(|(_, indices)| indices.as_slice())
2643    }
2644
2645    fn get_relation_key_indices(&self) -> Option<&[usize]> {
2646        self.relation_key_indices.as_deref()
2647    }
2648
2649    fn render_sink(
2650        &self,
2651        storage_state: &mut StorageState,
2652        sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
2653        sink_id: GlobalId,
2654        batches: SinkBatchStream<'scope>,
2655        key_is_synthetic: bool,
2656        _err_collection: VecCollection<'scope, Timestamp, DataflowError, Diff>,
2657    ) -> (
2658        StreamVec<'scope, Timestamp, HealthStatusMessage>,
2659        Vec<PressOnDropButton>,
2660    ) {
2661        let scope = batches.scope();
2662
2663        let write_handle = {
2664            let persist = Arc::clone(&storage_state.persist_clients);
2665            let shard_meta = sink.to_storage_metadata.clone();
2666            async move {
2667                let client = persist.open(shard_meta.persist_location).await?;
2668                let handle = client
2669                    .open_writer(
2670                        shard_meta.data_shard,
2671                        Arc::new(shard_meta.relation_desc),
2672                        Arc::new(UnitSchema),
2673                        Diagnostics::from_purpose("sink handle"),
2674                    )
2675                    .await?;
2676                Ok(handle)
2677            }
2678        };
2679
2680        let write_frontier = Rc::new(RefCell::new(Antichain::from_elem(Timestamp::minimum())));
2681        storage_state
2682            .sink_write_frontiers
2683            .insert(sink_id, Rc::clone(&write_frontier));
2684
2685        let (arrow_schema_with_ids, iceberg_schema) =
2686            match (|| -> Result<(ArrowSchema, Arc<Schema>), anyhow::Error> {
2687                let (arrow_schema_with_ids, iceberg_schema) =
2688                    relation_desc_to_iceberg_schema(&sink.from_desc)?;
2689
2690                Ok(if sink.envelope == SinkEnvelope::Append {
2691                    // For append mode, extend the Arrow and Iceberg schemas with the user-visible
2692                    // `_mz_diff` and `_mz_timestamp` columns. The minter uses `iceberg_schema` to create
2693                    // the Iceberg table, and `write_data_files` uses `arrow_schema_with_ids` when
2694                    // merging metadata. Both must include these columns before any operator starts.
2695                    let extended_arrow = build_schema_with_append_columns(&arrow_schema_with_ids);
2696                    let extended_iceberg = Arc::new(
2697                        arrow_schema_to_schema(&extended_arrow)
2698                            .context("Failed to build Iceberg schema with append columns")?,
2699                    );
2700                    (extended_arrow, extended_iceberg)
2701                } else {
2702                    (arrow_schema_with_ids, iceberg_schema)
2703                })
2704            })() {
2705                Ok(schemas) => schemas,
2706                Err(err) => {
2707                    let error_stream = std::iter::once(HealthStatusMessage {
2708                        id: None,
2709                        update: HealthStatusUpdate::halting(
2710                            format!("{}", err.display_with_causes()),
2711                            None,
2712                        ),
2713                        namespace: StatusNamespace::Iceberg,
2714                    })
2715                    .to_stream(scope);
2716                    return (error_stream, vec![]);
2717                }
2718            };
2719
2720        let metrics = Arc::new(
2721            storage_state
2722                .metrics
2723                .get_iceberg_sink_metrics(sink_id, scope.index()),
2724        );
2725
2726        let statistics = storage_state
2727            .aggregated_statistics
2728            .get_sink(&sink_id)
2729            .expect("statistics initialized")
2730            .clone();
2731
2732        let connection_for_minter = self.clone();
2733        let (batch_descriptions, table_ready, mint_status, mint_button) = mint_batch_descriptions(
2734            format!("{sink_id}-iceberg-mint"),
2735            sink_id,
2736            batches.clone(),
2737            sink,
2738            connection_for_minter,
2739            storage_state.storage_configuration.clone(),
2740            Arc::clone(&iceberg_schema),
2741        );
2742
2743        let connection_for_writer = self.clone();
2744        let (datafiles, write_status, write_button) = match sink.envelope {
2745            SinkEnvelope::Upsert => write_data_files::<UpsertEnvelopeHandler>(
2746                format!("{sink_id}-write-data-files"),
2747                batches,
2748                batch_descriptions.clone(),
2749                table_ready.clone(),
2750                sink_id,
2751                sink.from,
2752                key_is_synthetic,
2753                sink.as_of.clone(),
2754                connection_for_writer,
2755                storage_state.storage_configuration.clone(),
2756                Arc::new(arrow_schema_with_ids.clone()),
2757                Arc::clone(&metrics),
2758                statistics.clone(),
2759            ),
2760            SinkEnvelope::Append => write_data_files::<AppendEnvelopeHandler>(
2761                format!("{sink_id}-write-data-files"),
2762                batches,
2763                batch_descriptions.clone(),
2764                table_ready.clone(),
2765                sink_id,
2766                sink.from,
2767                key_is_synthetic,
2768                sink.as_of.clone(),
2769                connection_for_writer,
2770                storage_state.storage_configuration.clone(),
2771                Arc::new(arrow_schema_with_ids.clone()),
2772                Arc::clone(&metrics),
2773                statistics.clone(),
2774            ),
2775            SinkEnvelope::Debezium => {
2776                unreachable!("Iceberg sink only supports Upsert and Append envelopes")
2777            }
2778        };
2779
2780        let connection_for_committer = self.clone();
2781        let (commit_status, commit_button) = commit_to_iceberg(
2782            format!("{sink_id}-commit-to-iceberg"),
2783            sink_id,
2784            sink.version,
2785            datafiles,
2786            batch_descriptions,
2787            table_ready,
2788            Rc::clone(&write_frontier),
2789            connection_for_committer,
2790            storage_state.storage_configuration.clone(),
2791            write_handle,
2792            Arc::clone(&metrics),
2793            statistics,
2794        );
2795
2796        let running_status = Some(HealthStatusMessage {
2797            id: None,
2798            update: HealthStatusUpdate::running(),
2799            namespace: StatusNamespace::Iceberg,
2800        })
2801        .to_stream(scope);
2802
2803        let statuses =
2804            scope.concatenate([running_status, mint_status, write_status, commit_status]);
2805
2806        (statuses, vec![mint_button, write_button, commit_button])
2807    }
2808}