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