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