Skip to main content

mz_storage/sink/
iceberg.rs

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