1use 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
166const DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY: usize = 1024;
169const DEFAULT_ARRAY_BUILDER_DATA_CAPACITY: usize = 1024;
173
174const PARQUET_FILE_PREFIX: &str = "mz_data";
176const INITIAL_DESCRIPTIONS_TO_MINT: u64 = 3;
179
180struct WriterContext {
183 arrow_schema: Arc<ArrowSchema>,
185 current_schema: Arc<Schema>,
187 file_io: iceberg::io::FileIO,
189 location_generator: DefaultLocationGenerator,
191 file_name_generator: DefaultFileNameGenerator,
193 writer_properties: WriterProperties,
194}
195
196trait EnvelopeHandler: Send {
198 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 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 equality_ids: Vec<i32>,
221 pos_schema: Arc<Schema>,
223 eq_schema: Arc<Schema>,
225 eq_config: EqualityDeleteWriterConfig,
227 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 builder.with_max_seen_rows(0)
332 } else {
333 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 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 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 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 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
483fn 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
497fn 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
513fn 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
539fn merge_materialize_metadata_into_iceberg_schema(
544 materialize_arrow_schema: &ArrowSchema,
545 iceberg_schema: &Schema,
546) -> anyhow::Result<ArrowSchema> {
547 let iceberg_arrow_schema = schema_to_arrow_schema(iceberg_schema)
549 .context("Failed to convert Iceberg schema to Arrow schema")?;
550
551 let fields: Vec<Field> = iceberg_arrow_schema
553 .fields()
554 .iter()
555 .map(|iceberg_field| {
556 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
573fn merge_field_metadata_recursive(
575 iceberg_field: &Field,
576 mz_field: Option<&Field>,
577) -> anyhow::Result<Field> {
578 let mut metadata = iceberg_field.metadata().clone();
580
581 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 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 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
678fn 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
781async 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 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
915async 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 match catalog.load_table(&table_ident).await {
927 Ok(table) => {
928 Ok(table)
931 }
932 Err(err) => {
933 if matches!(err.kind(), ErrorKind::TableNotFound { .. })
934 || err
935 .message()
936 .contains("Tried to load a table that does not exist")
937 {
938 let table_creation = TableCreation::builder()
942 .name(table_name.clone())
943 .schema(schema.clone())
944 .build();
948
949 catalog
950 .create_table(&namespace_ident, table_creation)
951 .await
952 .with_context(|| {
953 format!(
954 "Failed to create Iceberg table '{}' in namespace '{}'",
955 table_name, namespace
956 )
957 })
958 } else {
959 Err(err).context("Failed to load Iceberg table")
961 }
962 }
963 }
964}
965
966fn retrieve_upper_from_snapshots(
971 snapshots: &mut [Arc<Snapshot>],
972) -> anyhow::Result<Option<(Antichain<Timestamp>, u64)>> {
973 snapshots.sort_by(|a, b| Ord::cmp(&b.sequence_number(), &a.sequence_number()));
974
975 for snapshot in snapshots {
976 let props = &snapshot.summary().additional_properties;
977 if let (Some(frontier_json), Some(sink_version_str)) =
978 (props.get("mz-frontier"), props.get("mz-sink-version"))
979 {
980 let frontier: Vec<Timestamp> = serde_json::from_str(frontier_json)
981 .context("Failed to deserialize frontier from snapshot properties")?;
982 let frontier = Antichain::from_iter(frontier);
983
984 let sink_version = sink_version_str
985 .parse::<u64>()
986 .context("Failed to parse mz-sink-version from snapshot properties")?;
987
988 return Ok(Some((frontier, sink_version)));
989 }
990 if snapshot.summary().operation.as_str() != "replace" {
991 anyhow::bail!(
996 "Iceberg table is in an inconsistent state: snapshot {} has operation '{}' but is missing 'mz-frontier' property. Schema or partition spec evolution is not supported.",
997 snapshot.snapshot_id(),
998 snapshot.summary().operation.as_str(),
999 );
1000 }
1001 }
1002
1003 Ok(None)
1004}
1005
1006fn relation_desc_to_iceberg_schema(
1016 desc: &mz_repr::RelationDesc,
1017) -> anyhow::Result<(ArrowSchema, SchemaRef)> {
1018 let arrow_schema =
1019 mz_arrow_util::builder::desc_to_schema_with_overrides(desc, iceberg_type_overrides)
1020 .context("Failed to convert RelationDesc to Iceberg-compatible Arrow schema")?;
1021
1022 let arrow_schema_with_ids = add_field_ids_to_arrow_schema(arrow_schema);
1023
1024 let iceberg_schema = arrow_schema_to_schema(&arrow_schema_with_ids)
1025 .context("Failed to convert Arrow schema to Iceberg schema")?;
1026
1027 Ok((arrow_schema_with_ids, Arc::new(iceberg_schema)))
1028}
1029
1030fn equality_ids_for_indices(
1035 current_schema: &Schema,
1036 materialize_arrow_schema: &ArrowSchema,
1037 equality_indices: &[usize],
1038) -> anyhow::Result<Vec<i32>> {
1039 let top_level_fields = current_schema.as_struct();
1040
1041 equality_indices
1042 .iter()
1043 .map(|index| {
1044 let mz_field = materialize_arrow_schema
1045 .fields()
1046 .get(*index)
1047 .with_context(|| format!("Equality delete key index {index} is out of bounds"))?;
1048 let field_name = mz_field.name();
1049 let iceberg_field = top_level_fields
1050 .field_by_name(field_name)
1051 .with_context(|| {
1052 format!(
1053 "Equality delete key column '{}' not found in Iceberg table schema",
1054 field_name
1055 )
1056 })?;
1057 Ok(iceberg_field.id)
1058 })
1059 .collect()
1060}
1061
1062fn build_schema_with_op_column(schema: &ArrowSchema) -> ArrowSchema {
1064 let mut fields: Vec<Arc<Field>> = schema.fields().iter().cloned().collect();
1065 fields.push(Arc::new(Field::new("__op", DataType::Int32, false)));
1066 ArrowSchema::new(fields)
1067}
1068
1069#[allow(clippy::disallowed_types)]
1074fn build_schema_with_append_columns(schema: &ArrowSchema) -> ArrowSchema {
1075 use mz_storage_types::sinks::{ICEBERG_APPEND_DIFF_COLUMN, ICEBERG_APPEND_TIMESTAMP_COLUMN};
1076 let mut fields: Vec<Arc<Field>> = schema.fields().iter().cloned().collect();
1077 fields.push(Arc::new(Field::new(
1078 ICEBERG_APPEND_DIFF_COLUMN,
1079 DataType::Int32,
1080 false,
1081 )));
1082 fields.push(Arc::new(Field::new(
1083 ICEBERG_APPEND_TIMESTAMP_COLUMN,
1084 DataType::Int64,
1085 false,
1086 )));
1087
1088 add_field_ids_to_arrow_schema(ArrowSchema::new(fields).with_metadata(schema.metadata().clone()))
1089}
1090
1091fn mint_batch_descriptions<'scope>(
1096 name: String,
1097 sink_id: GlobalId,
1098 input: SinkBatchStream<'scope>,
1099 sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
1100 connection: IcebergSinkConnection,
1101 storage_configuration: StorageConfiguration,
1102 initial_schema: SchemaRef,
1103) -> (
1104 StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
1105 StreamVec<'scope, Timestamp, Infallible>,
1106 StreamVec<'scope, Timestamp, HealthStatusMessage>,
1107 PressOnDropButton,
1108) {
1109 let scope = input.scope();
1110 let name_for_error = name.clone();
1111 let name_for_logging = name.clone();
1112 let mut builder = OperatorBuilder::new(name, scope.clone());
1113 let sink_version = sink.version;
1114
1115 let hashed_id = sink_id.hashed();
1116 let is_active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
1117 let (_, table_ready_stream) = builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
1118 let (batch_desc_output, batch_desc_stream) =
1119 builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
1120 let mut input = builder.new_input_for(input, Pipeline, &batch_desc_output);
1121
1122 let as_of = sink.as_of.clone();
1123 let commit_interval = sink
1124 .commit_interval
1125 .expect("the planner should have enforced this")
1126 .clone();
1127
1128 let (button, errors): (_, StreamVec<'scope, Timestamp, Rc<anyhow::Error>>) =
1129 builder.build_fallible(move |caps| {
1130 Box::pin(async move {
1131 let [table_ready_capset, capset]: &mut [_; 2] = caps.try_into().unwrap();
1132
1133 if !is_active_worker {
1134 return Ok(());
1136 }
1137
1138 let catalog = connection
1139 .catalog_connection
1140 .connect(&storage_configuration, InTask::Yes)
1141 .await
1142 .with_context(|| {
1143 format!(
1144 "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
1145 connection.catalog_connection.uri, connection.namespace, connection.table
1146 )
1147 })?;
1148
1149 let table = load_or_create_table(
1150 catalog.as_ref(),
1151 connection.namespace.clone(),
1152 connection.table.clone(),
1153 initial_schema.as_ref(),
1154 )
1155 .await?;
1156 debug!(
1157 ?sink_id,
1158 %name_for_logging,
1159 namespace = %connection.namespace,
1160 table = %connection.table,
1161 "iceberg mint loaded/created table"
1162 );
1163
1164 *table_ready_capset = CapabilitySet::new();
1165
1166 let mut snapshots: Vec<_> = table.metadata().snapshots().cloned().collect();
1167 let resume = retrieve_upper_from_snapshots(&mut snapshots)?;
1168 let (resume_upper, resume_version) = match resume {
1169 Some((f, v)) => (f, v),
1170 None => (Antichain::from_elem(Timestamp::minimum()), 0),
1171 };
1172 debug!(
1173 ?sink_id,
1174 %name_for_logging,
1175 resume_upper = %resume_upper.pretty(),
1176 resume_version,
1177 as_of = %as_of.pretty(),
1178 "iceberg mint resume position loaded"
1179 );
1180
1181 let overcompacted =
1183 *resume_upper != [Timestamp::minimum()] &&
1185 PartialOrder::less_than(&resume_upper, &as_of);
1187
1188 if overcompacted {
1189 let err = format!(
1190 "{name_for_error}: input compacted past resume upper: as_of {}, resume_upper: {}",
1191 as_of.pretty(),
1192 resume_upper.pretty()
1193 );
1194 return Err(anyhow::anyhow!("{err}"));
1198 };
1199
1200 if resume_version > sink_version {
1201 anyhow::bail!("Fenced off by newer sink version: resume_version {}, sink_version {}", resume_version, sink_version);
1202 }
1203
1204 let mut initialized = false;
1205 let mut observed_frontier;
1206 let mut minted_batches = VecDeque::new();
1211
1212 let catchup_start = if *resume_upper == [Timestamp::minimum()] {
1215 let batch_upper = Antichain::from_elem(
1217 as_of.as_option().expect("as_of not empty").step_forward());
1218 let batch = (as_of.clone(), batch_upper.clone());
1219 minted_batches.push_back(batch.clone());
1220 batch_desc_output.give(&capset[0], batch);
1221 capset.downgrade(batch_upper.clone());
1222
1223 batch_upper
1225 } else {
1226 resume_upper.clone()
1228 };
1229
1230 loop {
1231 if let Some(event) = input.next().await {
1232 match event {
1233 Event::Data(_, _) => continue,
1234 Event::Progress(frontier) => {
1235 observed_frontier = frontier;
1236 }
1237 }
1238 } else {
1239 return Ok(());
1240 }
1241
1242 if !initialized {
1243 if observed_frontier.is_empty() {
1244 if catchup_start.is_empty() {
1251 return Ok(());
1254 }
1255 debug!(
1256 ?sink_id,
1257 %name_for_logging,
1258 batch_lower = %catchup_start.pretty(),
1259 "iceberg mint input closed before initialization; minting final batch"
1260 );
1261 let batch = (catchup_start.clone(), Antichain::new());
1262 batch_desc_output.give(&capset[0], batch);
1263 return Ok(());
1264 }
1265
1266 if !PartialOrder::less_than(&catchup_start, &observed_frontier)
1269 {
1270 continue;
1271 }
1272
1273 let mut batch_descriptions = vec![];
1274 let mut current_upper = observed_frontier.clone();
1275 let current_upper_ts = observed_frontier.as_option().expect("frontier not empty").clone();
1276 debug!(
1277 ?sink_id,
1278 %name_for_logging,
1279 batch_lower = %catchup_start.pretty(),
1280 current_upper = %current_upper.pretty(),
1281 "iceberg mint initializing (catch-up batch)"
1282 );
1283 debug!(
1284 "{}: creating catch-up batch [{}, {})",
1285 name_for_logging,
1286 catchup_start.pretty(),
1287 current_upper.pretty()
1288 );
1289 batch_descriptions.push((catchup_start.clone(), current_upper.clone()));
1290
1291 for i in 1..INITIAL_DESCRIPTIONS_TO_MINT {
1293 let duration_millis = commit_interval.as_millis()
1294 .checked_mul(u128::from(i))
1295 .expect("commit interval multiplication overflow");
1296 let duration_ts = Timestamp::new(
1297 u64::try_from(duration_millis)
1298 .expect("commit interval too large for u64"),
1299 );
1300 let desired_batch_upper = Antichain::from_elem(
1301 current_upper_ts.step_forward_by(&duration_ts),
1302 );
1303
1304 let batch_description =
1305 (current_upper.clone(), desired_batch_upper.clone());
1306 debug!(
1307 "{}: minting future batch {}/{} [{}, {})",
1308 name_for_logging,
1309 i,
1310 INITIAL_DESCRIPTIONS_TO_MINT,
1311 current_upper.pretty(),
1312 desired_batch_upper.pretty()
1313 );
1314 current_upper = batch_description.1.clone();
1315 batch_descriptions.push(batch_description);
1316 }
1317
1318 minted_batches.extend(batch_descriptions.clone());
1319
1320 for desc in batch_descriptions {
1321 batch_desc_output.give(&capset[0], desc);
1322 }
1323
1324 capset.downgrade(current_upper);
1325
1326 initialized = true;
1327 } else {
1328 if observed_frontier.is_empty() {
1329 return Ok(());
1331 }
1332 while let Some(oldest_desc) = minted_batches.front() {
1335 let oldest_upper = &oldest_desc.1;
1336 if !PartialOrder::less_equal(oldest_upper, &observed_frontier) {
1337 break;
1338 }
1339
1340 let newest_upper = minted_batches.back().unwrap().1.clone();
1341 let new_lower = newest_upper.clone();
1342 let duration_ts = Timestamp::new(commit_interval.as_millis()
1343 .try_into()
1344 .expect("commit interval too large for u64"));
1345 let new_upper = Antichain::from_elem(newest_upper
1346 .as_option()
1347 .unwrap()
1348 .step_forward_by(&duration_ts));
1349
1350 let new_batch_description = (new_lower.clone(), new_upper.clone());
1351 minted_batches.pop_front();
1352 minted_batches.push_back(new_batch_description.clone());
1353
1354 batch_desc_output.give(&capset[0], new_batch_description);
1355
1356 capset.downgrade(new_upper);
1357 }
1358 }
1359 }
1360 })
1361 });
1362
1363 let statuses = errors.map(|error| HealthStatusMessage {
1364 id: None,
1365 update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
1366 namespace: StatusNamespace::Iceberg,
1367 });
1368 (
1369 batch_desc_stream,
1370 table_ready_stream,
1371 statuses,
1372 button.press_on_drop(),
1373 )
1374}
1375
1376#[derive(Clone, Debug, Serialize, Deserialize)]
1377#[serde(try_from = "AvroDataFile", into = "AvroDataFile")]
1378struct SerializableDataFile {
1379 pub data_file: DataFile,
1380 pub schema: Schema,
1381}
1382
1383#[derive(Clone, Debug, Serialize, Deserialize)]
1391struct AvroDataFile {
1392 pub data_file: Vec<u8>,
1393 pub schema: Vec<u8>,
1395}
1396
1397impl From<SerializableDataFile> for AvroDataFile {
1398 fn from(value: SerializableDataFile) -> Self {
1399 let mut data_file = Vec::new();
1400 write_data_files_to_avro(
1401 &mut data_file,
1402 [value.data_file],
1403 &StructType::new(vec![]),
1404 FormatVersion::V2,
1405 )
1406 .expect("serialization into buffer");
1407 let schema = serde_json::to_vec(&value.schema).expect("schema serialization");
1408 AvroDataFile { data_file, schema }
1409 }
1410}
1411
1412impl TryFrom<AvroDataFile> for SerializableDataFile {
1413 type Error = String;
1414
1415 fn try_from(value: AvroDataFile) -> Result<Self, Self::Error> {
1416 let schema: Schema = serde_json::from_slice(&value.schema)
1417 .map_err(|e| format!("Failed to deserialize schema: {}", e))?;
1418 let data_files = read_data_files_from_avro(
1419 &mut &*value.data_file,
1420 &schema,
1421 0,
1422 &StructType::new(vec![]),
1423 FormatVersion::V2,
1424 )
1425 .map_err_to_string_with_causes()?;
1426 let Some(data_file) = data_files.into_iter().next() else {
1427 return Err("No DataFile found in Avro data".into());
1428 };
1429 Ok(SerializableDataFile { data_file, schema })
1430 }
1431}
1432
1433#[derive(Clone, Debug, Serialize, Deserialize)]
1435struct BoundedDataFile {
1436 pub data_file: SerializableDataFile,
1437 pub batch_desc: (Antichain<Timestamp>, Antichain<Timestamp>),
1438}
1439
1440impl BoundedDataFile {
1441 pub fn new(
1442 file: DataFile,
1443 schema: Schema,
1444 batch_desc: (Antichain<Timestamp>, Antichain<Timestamp>),
1445 ) -> Self {
1446 Self {
1447 data_file: SerializableDataFile {
1448 data_file: file,
1449 schema,
1450 },
1451 batch_desc,
1452 }
1453 }
1454
1455 pub fn batch_desc(&self) -> &(Antichain<Timestamp>, Antichain<Timestamp>) {
1456 &self.batch_desc
1457 }
1458
1459 pub fn data_file(&self) -> &DataFile {
1460 &self.data_file.data_file
1461 }
1462
1463 pub fn into_data_file(self) -> DataFile {
1464 self.data_file.data_file
1465 }
1466}
1467
1468#[derive(Clone, Debug, Default)]
1470struct BoundedDataFileSet {
1471 pub data_files: Vec<BoundedDataFile>,
1472}
1473
1474fn write_data_files<'scope, H: EnvelopeHandler + 'static>(
1480 name: String,
1481 input: SinkBatchStream<'scope>,
1482 batch_desc_input: StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
1483 table_ready_stream: StreamVec<'scope, Timestamp, Infallible>,
1484 sink_id: GlobalId,
1485 from_id: GlobalId,
1486 key_is_synthetic: bool,
1487 as_of: Antichain<Timestamp>,
1488 connection: IcebergSinkConnection,
1489 storage_configuration: StorageConfiguration,
1490 materialize_arrow_schema: Arc<ArrowSchema>,
1491 metrics: Arc<IcebergSinkMetrics>,
1492 statistics: SinkStatistics,
1493) -> (
1494 StreamVec<'scope, Timestamp, BoundedDataFile>,
1495 StreamVec<'scope, Timestamp, HealthStatusMessage>,
1496 PressOnDropButton,
1497) {
1498 let scope = input.scope();
1499 let name_for_logging = name.clone();
1500 let mut builder = OperatorBuilder::new(name, scope.clone());
1501
1502 let (output, output_stream) = builder.new_output::<CapacityContainerBuilder<_>>();
1503
1504 let mut table_ready_input = builder.new_disconnected_input(table_ready_stream, Pipeline);
1505 let mut batch_desc_input =
1506 builder.new_input_for(batch_desc_input.broadcast(), Pipeline, &output);
1507 let mut input = builder.new_disconnected_input(input, Pipeline);
1508
1509 let (button, errors): (_, StreamVec<'scope, Timestamp, Rc<anyhow::Error>>) = builder
1510 .build_fallible(move |caps| {
1511 Box::pin(async move {
1512 let [capset]: &mut [_; 1] = caps.try_into().unwrap();
1513 let catalog = connection
1514 .catalog_connection
1515 .connect(&storage_configuration, InTask::Yes)
1516 .await
1517 .with_context(|| {
1518 format!(
1519 "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
1520 connection.catalog_connection.uri,
1521 connection.namespace,
1522 connection.table
1523 )
1524 })?;
1525
1526 let namespace_ident = NamespaceIdent::new(connection.namespace.clone());
1527 let table_ident = TableIdent::new(namespace_ident, connection.table.clone());
1528 while let Some(_) = table_ready_input.next().await {
1529 }
1531 let table = catalog.load_table(&table_ident).await.with_context(|| {
1532 format!(
1533 "Failed to load Iceberg table '{}.{}' in write_data_files operator",
1534 connection.namespace, connection.table
1535 )
1536 })?;
1537
1538 let table_metadata = table.metadata().clone();
1539 let current_schema = Arc::clone(table_metadata.current_schema());
1540
1541 let arrow_schema = Arc::new(
1545 merge_materialize_metadata_into_iceberg_schema(
1546 materialize_arrow_schema.as_ref(),
1547 current_schema.as_ref(),
1548 )
1549 .context("Failed to merge Materialize metadata into Iceberg schema")?,
1550 );
1551
1552 let location = table_metadata.location();
1556 let corrected_location = match location.rsplit_once("/metadata/") {
1557 Some((a, b)) if b.ends_with(".metadata.json") => a,
1558 _ => location,
1559 };
1560
1561 let data_location = format!("{}/data", corrected_location);
1562 let location_generator =
1563 DefaultLocationGenerator::with_data_location(data_location);
1564
1565 let unique_suffix = format!("-{}", uuid::Uuid::new_v4());
1567 let file_name_generator = DefaultFileNameGenerator::new(
1568 PARQUET_FILE_PREFIX.to_string(),
1569 Some(unique_suffix),
1570 iceberg::spec::DataFileFormat::Parquet,
1571 );
1572
1573 let file_io = table.file_io().clone();
1574
1575 let writer_properties = WriterProperties::new();
1576
1577 let ctx = WriterContext {
1578 arrow_schema,
1579 current_schema: Arc::clone(¤t_schema),
1580 file_io,
1581 location_generator,
1582 file_name_generator,
1583 writer_properties,
1584 };
1585 let handler = H::new(ctx, &connection, &materialize_arrow_schema)?;
1586 let mut pk_warner =
1587 (!key_is_synthetic).then(|| PkViolationWarner::new(sink_id, from_id));
1588
1589 let mut stashed_rows: VecDeque<ArcBatch<OrdValBatch<_>>> = VecDeque::new();
1593
1594 let mut in_flight_batches: VecDeque<(
1598 (Antichain<Timestamp>, Antichain<Timestamp>),
1599 Box<dyn IcebergWriter>,
1600 )> = VecDeque::new();
1601
1602 let mut last_batch_desc: Option<BatchDescription> = None;
1606 let mut last_input_bounds: Option<(Antichain<Timestamp>, Antichain<Timestamp>)> =
1607 None;
1608
1609 let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
1610 let mut input_frontier = Antichain::from_elem(Timestamp::minimum());
1611
1612 while !(batch_description_frontier.is_empty() && input_frontier.is_empty()) {
1613 tokio::select! {
1614 _ = batch_desc_input.ready() => {},
1615 _ = input.ready() => {}
1616 }
1617
1618 while let Some(event) = batch_desc_input.next_sync() {
1622 match event {
1623 Event::Data(_cap, data) => {
1624 for batch_desc in data {
1625 let (lower, upper) = &batch_desc;
1626
1627 if let Some((prev_lower, prev_upper)) = last_batch_desc.as_ref()
1628 {
1629 if prev_upper != lower {
1630 anyhow::bail!(
1631 "batch descriptions must arrive in order, non-overlapping, \
1632 and without gaps: previous [{}, {}), new [{}, {})",
1633 prev_lower.pretty(),
1634 prev_upper.pretty(),
1635 lower.pretty(),
1636 upper.pretty(),
1637 );
1638 }
1639 }
1640 last_batch_desc = Some(batch_desc.clone());
1641
1642 let is_snapshot = lower == &as_of;
1644 debug!(
1645 "{}: received batch description [{}, {}), snapshot={}",
1646 name_for_logging,
1647 lower.pretty(),
1648 upper.pretty(),
1649 is_snapshot
1650 );
1651 let batch_writer = handler.create_writer(is_snapshot).await?;
1652 in_flight_batches.push_back((batch_desc.clone(), batch_writer));
1653 }
1654 }
1655 Event::Progress(frontier) => {
1656 batch_description_frontier = frontier;
1657 }
1658 }
1659 }
1660
1661 while let Some(event) = input.next_sync() {
1663 match event {
1664 Event::Data(_cap, data) => {
1665 for rows in &data {
1666 if let Some((prev_lower, prev_upper)) =
1667 last_input_bounds.as_ref()
1668 {
1669 if !PartialOrder::less_equal(prev_upper, rows.lower()) {
1673 anyhow::bail!(
1674 "input batches must arrive in order and \
1675 non-overlapping: previous [{}, {}), new [{}, {})",
1676 prev_lower.pretty(),
1677 prev_upper.pretty(),
1678 rows.lower().pretty(),
1679 rows.upper().pretty(),
1680 );
1681 }
1682 }
1683 last_input_bounds =
1684 Some((rows.lower().clone(), rows.upper().clone()));
1685
1686 stashed_rows.push_back(rows.clone());
1687 }
1688 }
1689 Event::Progress(frontier) => {
1690 input_frontier = frontier;
1691 }
1692 }
1693 }
1694
1695 metrics.stashed_rows.set(u64::cast_from(
1696 stashed_rows.iter().map(|rows| rows.len()).sum::<usize>(),
1697 ));
1698
1699 let mut staged_messages_since_flush: u64 = 0;
1704
1705 let write_rows = async |rows: &OrdValBatch<_>,
1707 (lower, upper): BatchDescription,
1708 batch_writer: &mut Box<dyn IcebergWriter>|
1709 -> Result<(), anyhow::Error> {
1710 for_each_diff_pair_async(
1711 rows,
1712 Some(lower),
1713 Some(upper),
1714 async |key, time, diff_pair| -> Result<(), anyhow::Error> {
1715 if let Some(warner) = pk_warner.as_mut() {
1716 warner.observe(key, time);
1717 }
1718
1719 let record_batch = handler
1720 .row_to_batch(diff_pair, time)
1721 .context("failed to convert row to recordbatch")?;
1722 staged_messages_since_flush +=
1723 u64::cast_from(record_batch.num_rows());
1724 batch_writer
1725 .write(record_batch)
1726 .await
1727 .context("failed to write recordbatch")?;
1728 if staged_messages_since_flush >= 10_000 {
1729 statistics.inc_messages_staged_by(staged_messages_since_flush);
1730 staged_messages_since_flush = 0;
1731 }
1732 Ok(())
1733 },
1734 )
1735 .await?;
1736 if let Some(warner) = pk_warner.as_mut() {
1740 warner.flush();
1741 }
1742 Ok(())
1743 };
1744
1745 let close_batch = async |batch_desc: BatchDescription,
1747 batch_writer: &mut Box<dyn IcebergWriter>|
1748 -> Result<(), anyhow::Error> {
1749 let close_started_at = Instant::now();
1750 let data_files = batch_writer.close().await;
1751 metrics
1752 .writer_close_duration_seconds
1753 .observe(close_started_at.elapsed().as_secs_f64());
1754 let data_files = data_files.context("Failed to close batch writer")?;
1755 debug!(
1756 "{}: closed batch [{}, {}), wrote {} files",
1757 name_for_logging,
1758 batch_desc.0.pretty(),
1759 batch_desc.1.pretty(),
1760 data_files.len()
1761 );
1762 for data_file in data_files {
1763 match data_file.content_type() {
1764 iceberg::spec::DataContentType::Data => {
1765 metrics.data_files_written.inc();
1766 }
1767 iceberg::spec::DataContentType::PositionDeletes
1768 | iceberg::spec::DataContentType::EqualityDeletes => {
1769 metrics.delete_files_written.inc();
1770 }
1771 }
1772 statistics.inc_bytes_staged_by(data_file.file_size_in_bytes());
1773 let file = BoundedDataFile::new(
1774 data_file,
1775 current_schema.as_ref().clone(),
1776 batch_desc.clone(),
1777 );
1778 output.give(&capset[0], file);
1779 }
1780
1781 capset.downgrade(batch_desc.1.clone());
1784 Ok(())
1785 };
1786
1787 with_ready_batches(
1789 input_frontier.clone(),
1790 &mut stashed_rows,
1791 batch_description_frontier.clone(),
1792 &mut in_flight_batches,
1793 write_rows,
1794 close_batch,
1795 )
1796 .await?;
1797
1798 if staged_messages_since_flush > 0 {
1799 statistics.inc_messages_staged_by(staged_messages_since_flush);
1800 }
1801 metrics.stashed_rows.set(u64::cast_from(
1802 stashed_rows.iter().map(|rows| rows.len()).sum::<usize>(),
1803 ));
1804 }
1805 Ok(())
1806 })
1807 });
1808
1809 let statuses = errors.map(|error| HealthStatusMessage {
1810 id: None,
1811 update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
1812 namespace: StatusNamespace::Iceberg,
1813 });
1814 (output_stream, statuses, button.press_on_drop())
1815}
1816
1817type BatchDescription = (Antichain<Timestamp>, Antichain<Timestamp>);
1819
1820async fn with_ready_batches<L: Layout, W, Write, Close>(
1832 input_frontier: Antichain<Timestamp>,
1833 input_batches: &mut VecDeque<ArcBatch<OrdValBatch<L>>>,
1834 output_frontier: Antichain<Timestamp>,
1835 output_batches: &mut VecDeque<(BatchDescription, W)>,
1836 mut write_rows: Write,
1837 mut close_batch: Close,
1838) -> Result<(), anyhow::Error>
1839where
1840 L::TimeContainer: BatchContainer<Owned = Timestamp>,
1841 Write: AsyncFnMut(&OrdValBatch<L>, BatchDescription, &mut W) -> Result<(), anyhow::Error>,
1842 Close: AsyncFnMut(BatchDescription, &mut W) -> Result<(), anyhow::Error>,
1843{
1844 loop {
1845 {
1846 let output_lower = output_batches
1849 .front()
1850 .map_or(&output_frontier, |((lower, _), _)| lower);
1851 while input_batches
1852 .pop_front_if(|rows| PartialOrder::less_equal(rows.upper(), output_lower))
1853 .is_some()
1854 {}
1855 }
1856
1857 {
1858 let input_lower = input_batches
1861 .front()
1862 .map_or(&input_frontier, |rows| rows.lower());
1863 while let Some((batch_desc, mut batch_writer)) =
1864 output_batches.pop_front_if(|((_, batch_upper), _)| {
1865 PartialOrder::less_equal(batch_upper, input_lower)
1866 })
1867 {
1868 close_batch(batch_desc, &mut batch_writer).await?;
1869 }
1870 }
1871
1872 let Some((batch_desc, batch_writer)) = output_batches.front_mut() else {
1873 break;
1875 };
1876
1877 let Some(rows) = input_batches.front() else {
1878 break;
1880 };
1881
1882 write_rows(rows, batch_desc.clone(), batch_writer).await?;
1890 let output_upper = batch_desc.1.clone();
1891 let rows_upper = rows.upper();
1892 if PartialOrder::less_equal(&output_upper, rows_upper) {
1893 let (batch_desc, mut batch_writer) =
1895 output_batches.pop_front().expect("already checked front");
1896 close_batch(batch_desc, &mut batch_writer).await?;
1897 }
1898 if PartialOrder::less_equal(rows_upper, &output_upper) {
1899 input_batches.pop_front();
1901 }
1902
1903 }
1907
1908 Ok(())
1909}
1910
1911#[cfg(test)]
1912mod tests {
1913 use super::*;
1914 use iceberg::spec::{PrimitiveType, Type};
1915 use mz_repr::SqlScalarType;
1916 use mz_storage_types::sinks::ICEBERG_UINT64_DECIMAL_PRECISION;
1917
1918 #[mz_ore::test]
1919 fn test_iceberg_type_overrides() {
1920 let result = iceberg_type_overrides(&SqlScalarType::UInt16);
1922 assert_eq!(result.unwrap().0, DataType::Int32);
1923
1924 let result = iceberg_type_overrides(&SqlScalarType::UInt32);
1926 assert_eq!(result.unwrap().0, DataType::Int64);
1927
1928 let result = iceberg_type_overrides(&SqlScalarType::UInt64);
1930 assert_eq!(
1931 result.unwrap().0,
1932 DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
1933 );
1934
1935 let result = iceberg_type_overrides(&SqlScalarType::MzTimestamp);
1937 assert_eq!(
1938 result.unwrap().0,
1939 DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
1940 );
1941
1942 assert!(iceberg_type_overrides(&SqlScalarType::Int32).is_none());
1944 assert!(iceberg_type_overrides(&SqlScalarType::String).is_none());
1945 assert!(iceberg_type_overrides(&SqlScalarType::Bool).is_none());
1946 }
1947
1948 #[mz_ore::test]
1949 fn test_iceberg_schema_with_nested_uint64() {
1950 let desc = mz_repr::RelationDesc::builder()
1953 .with_column(
1954 "items",
1955 SqlScalarType::List {
1956 element_type: Box::new(SqlScalarType::UInt64),
1957 custom_id: None,
1958 }
1959 .nullable(true),
1960 )
1961 .finish();
1962
1963 let schema =
1964 mz_arrow_util::builder::desc_to_schema_with_overrides(&desc, iceberg_type_overrides)
1965 .expect("schema conversion should succeed");
1966
1967 if let DataType::List(field) = schema.field(0).data_type() {
1969 assert_eq!(
1970 field.data_type(),
1971 &DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
1972 );
1973 } else {
1974 panic!("Expected List type");
1975 }
1976 }
1977
1978 #[mz_ore::test]
1979 fn test_iceberg_interval_override() {
1980 let result = iceberg_type_overrides(&SqlScalarType::Interval);
1982 assert_eq!(result.unwrap().0, DataType::LargeUtf8);
1983
1984 let desc = mz_repr::RelationDesc::builder()
1986 .with_column("id", SqlScalarType::Int32.nullable(false))
1987 .with_column("dur", SqlScalarType::Interval.nullable(true))
1988 .finish();
1989
1990 let (arrow_schema, iceberg_schema) =
1991 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
1992
1993 assert_eq!(arrow_schema.field(1).data_type(), &DataType::LargeUtf8);
1995
1996 let field = iceberg_schema
1998 .as_struct()
1999 .field_by_name("dur")
2000 .expect("field should exist");
2001 assert_eq!(*field.field_type, Type::Primitive(PrimitiveType::String));
2002 }
2003
2004 #[mz_ore::test]
2005 fn test_iceberg_range_schema() {
2006 let desc = mz_repr::RelationDesc::builder()
2008 .with_column("id", SqlScalarType::Int32.nullable(false))
2009 .with_column(
2010 "r",
2011 SqlScalarType::Range {
2012 element_type: Box::new(SqlScalarType::Int32),
2013 }
2014 .nullable(true),
2015 )
2016 .finish();
2017
2018 let (_arrow_schema, iceberg_schema) =
2019 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2020
2021 let field = iceberg_schema
2023 .as_struct()
2024 .field_by_name("r")
2025 .expect("field should exist");
2026 assert!(
2027 matches!(&*field.field_type, Type::Struct(_)),
2028 "range should be struct, got: {:?}",
2029 field.field_type
2030 );
2031 }
2032
2033 #[mz_ore::test]
2034 fn equality_ids_follow_iceberg_field_ids() {
2035 let map_entries = Field::new(
2036 "entries",
2037 DataType::Struct(
2038 vec![
2039 Field::new("key", DataType::Utf8, false),
2040 Field::new("value", DataType::Utf8, true),
2041 ]
2042 .into(),
2043 ),
2044 false,
2045 );
2046 let materialize_arrow_schema = ArrowSchema::new(vec![
2047 Field::new("attrs", DataType::Map(Arc::new(map_entries), false), true),
2048 Field::new("key_col", DataType::Int32, false),
2049 ]);
2050 let materialize_arrow_schema = add_field_ids_to_arrow_schema(materialize_arrow_schema);
2051 let iceberg_schema = arrow_schema_to_schema(&materialize_arrow_schema)
2052 .expect("schema conversion should succeed");
2053
2054 let equality_ids =
2055 equality_ids_for_indices(&iceberg_schema, &materialize_arrow_schema, &[1])
2056 .expect("field lookup should succeed");
2057
2058 let expected_id = iceberg_schema
2059 .as_struct()
2060 .field_by_name("key_col")
2061 .expect("top-level field should exist")
2062 .id;
2063 assert_eq!(equality_ids, vec![expected_id]);
2064 assert_ne!(expected_id, 2);
2065 }
2066
2067 #[mz_ore::test]
2072 #[allow(clippy::disallowed_types)]
2073 fn merge_map_entries_preserves_value_extension_metadata() {
2074 use std::collections::HashMap;
2075
2076 let mz_value_metadata = HashMap::from([(
2077 ARROW_EXTENSION_NAME_KEY.to_string(),
2078 "materialize.v1.string".to_string(),
2079 )]);
2080 let mz_entries = Field::new(
2081 "entries",
2082 DataType::Struct(
2083 vec![
2084 Field::new("keys", DataType::Utf8, false),
2085 Field::new("values", DataType::Utf8, true).with_metadata(mz_value_metadata),
2086 ]
2087 .into(),
2088 ),
2089 false,
2090 );
2091 let mz_map = Field::new("m", DataType::Map(Arc::new(mz_entries), false), true)
2092 .with_metadata(HashMap::from([(
2093 ARROW_EXTENSION_NAME_KEY.to_string(),
2094 "materialize.v1.map".to_string(),
2095 )]));
2096
2097 let iceberg_entries = Field::new(
2098 "key_value",
2099 DataType::Struct(
2100 vec![
2101 Field::new("key", DataType::Utf8, false),
2102 Field::new("value", DataType::Utf8, true),
2103 ]
2104 .into(),
2105 ),
2106 false,
2107 );
2108 let iceberg_map = Field::new("m", DataType::Map(Arc::new(iceberg_entries), false), true);
2109
2110 let merged = merge_field_metadata_recursive(&iceberg_map, Some(&mz_map))
2111 .expect("merge should succeed");
2112
2113 let entries = match merged.data_type() {
2114 DataType::Map(entries, _) => entries.as_ref(),
2115 other => panic!("expected Map, got {other:?}"),
2116 };
2117 let entry_fields = match entries.data_type() {
2118 DataType::Struct(fields) => fields,
2119 other => panic!("expected Struct, got {other:?}"),
2120 };
2121 assert_eq!(entry_fields[0].name(), "key");
2123 assert_eq!(entry_fields[1].name(), "value");
2124 assert_eq!(
2127 entry_fields[1].metadata().get(ARROW_EXTENSION_NAME_KEY),
2128 Some(&"materialize.v1.string".to_string()),
2129 );
2130 }
2131
2132 mod with_ready_batches {
2133 use differential_dataflow::trace::Batch;
2134 use differential_dataflow::trace::implementations::Vector;
2135
2136 use super::*;
2137
2138 type TestBatch = OrdValBatch<Vector<((u64, u64), Timestamp, Diff)>>;
2139
2140 fn frontier(t: Option<u64>) -> Antichain<Timestamp> {
2142 t.map_or_else(Antichain::new, |t| Antichain::from_elem(Timestamp::new(t)))
2143 }
2144
2145 fn span(lower: u64, upper: Option<u64>) -> BatchDescription {
2147 (frontier(Some(lower)), frontier(upper))
2148 }
2149
2150 fn input(lower: u64, upper: Option<u64>) -> ArcBatch<TestBatch> {
2153 let (lower, upper) = span(lower, upper);
2154 ArcBatch(Arc::new(TestBatch::empty(lower, upper)))
2155 }
2156
2157 #[derive(Debug, PartialEq)]
2158 enum Call {
2159 Write(BatchDescription, BatchDescription),
2161 Close(BatchDescription),
2162 }
2163
2164 async fn run(
2167 input_frontier: Antichain<Timestamp>,
2168 input_batches: &mut VecDeque<ArcBatch<TestBatch>>,
2169 output_frontier: Antichain<Timestamp>,
2170 output_batches: &mut VecDeque<(BatchDescription, ())>,
2171 ) -> Vec<Call> {
2172 let calls = RefCell::new(vec![]);
2173 with_ready_batches(
2174 input_frontier,
2175 input_batches,
2176 output_frontier,
2177 output_batches,
2178 async |rows: &TestBatch, desc, _writer: &mut ()| {
2179 let bounds = (rows.lower().clone(), rows.upper().clone());
2180 calls.borrow_mut().push(Call::Write(bounds, desc));
2181 Ok(())
2182 },
2183 async |desc, _writer: &mut ()| {
2184 calls.borrow_mut().push(Call::Close(desc));
2185 Ok(())
2186 },
2187 )
2188 .await
2189 .expect("test callbacks never fail");
2190 calls.into_inner()
2191 }
2192
2193 #[mz_ore::test(tokio::test)]
2194 async fn input_batch_spanning_multiple_output_batches() {
2195 let mut inputs = VecDeque::from([input(0, Some(30))]);
2196 let mut outputs = VecDeque::from([
2197 (span(0, Some(10)), ()),
2198 (span(10, Some(20)), ()),
2199 (span(20, Some(30)), ()),
2200 ]);
2201
2202 let calls = run(
2203 frontier(Some(30)),
2204 &mut inputs,
2205 frontier(Some(30)),
2206 &mut outputs,
2207 )
2208 .await;
2209
2210 assert_eq!(
2213 calls,
2214 vec![
2215 Call::Write(span(0, Some(30)), span(0, Some(10))),
2216 Call::Close(span(0, Some(10))),
2217 Call::Write(span(0, Some(30)), span(10, Some(20))),
2218 Call::Close(span(10, Some(20))),
2219 Call::Write(span(0, Some(30)), span(20, Some(30))),
2220 Call::Close(span(20, Some(30))),
2221 ]
2222 );
2223 assert!(inputs.is_empty());
2224 assert!(outputs.is_empty());
2225 }
2226
2227 #[mz_ore::test(tokio::test)]
2228 async fn output_batch_spanning_multiple_input_batches() {
2229 let mut inputs =
2230 VecDeque::from([input(0, Some(10)), input(10, Some(20)), input(20, Some(30))]);
2231 let mut outputs = VecDeque::from([(span(0, Some(30)), ())]);
2232
2233 let calls = run(
2234 frontier(Some(30)),
2235 &mut inputs,
2236 frontier(Some(30)),
2237 &mut outputs,
2238 )
2239 .await;
2240
2241 assert_eq!(
2242 calls,
2243 vec![
2244 Call::Write(span(0, Some(10)), span(0, Some(30))),
2245 Call::Write(span(10, Some(20)), span(0, Some(30))),
2246 Call::Write(span(20, Some(30)), span(0, Some(30))),
2247 Call::Close(span(0, Some(30))),
2248 ]
2249 );
2250 assert!(inputs.is_empty());
2251 assert!(outputs.is_empty());
2252 }
2253
2254 #[mz_ore::test(tokio::test)]
2255 async fn input_batch_retained_for_future_output_batches() {
2256 let mut inputs = VecDeque::from([input(0, Some(30))]);
2257 let mut outputs = VecDeque::from([(span(0, Some(10)), ())]);
2258
2259 let calls = run(
2260 frontier(Some(30)),
2261 &mut inputs,
2262 frontier(Some(10)),
2263 &mut outputs,
2264 )
2265 .await;
2266
2267 assert_eq!(
2270 calls,
2271 vec![
2272 Call::Write(span(0, Some(30)), span(0, Some(10))),
2273 Call::Close(span(0, Some(10))),
2274 ]
2275 );
2276 assert_eq!(inputs.len(), 1);
2277 assert!(outputs.is_empty());
2278 }
2279
2280 #[mz_ore::test(tokio::test)]
2281 async fn already_committed_input_batches_dropped_unwritten() {
2282 let mut inputs = VecDeque::from([input(0, Some(10)), input(10, Some(20))]);
2283 let mut outputs = VecDeque::from([(span(20, Some(30)), ())]);
2284
2285 let calls = run(
2286 frontier(Some(20)),
2287 &mut inputs,
2288 frontier(Some(30)),
2289 &mut outputs,
2290 )
2291 .await;
2292
2293 assert_eq!(calls, vec![]);
2297 assert!(inputs.is_empty());
2298 assert_eq!(outputs.len(), 1);
2299 }
2300
2301 #[mz_ore::test(tokio::test)]
2302 async fn output_batch_closes_empty_once_input_frontier_passes() {
2303 let mut outputs = VecDeque::from([(span(0, Some(10)), ())]);
2304
2305 let calls = run(
2308 frontier(Some(5)),
2309 &mut VecDeque::new(),
2310 frontier(Some(10)),
2311 &mut outputs,
2312 )
2313 .await;
2314 assert_eq!(calls, vec![]);
2315 assert_eq!(outputs.len(), 1);
2316
2317 let calls = run(
2320 frontier(Some(10)),
2321 &mut VecDeque::new(),
2322 frontier(Some(10)),
2323 &mut outputs,
2324 )
2325 .await;
2326 assert_eq!(calls, vec![Call::Close(span(0, Some(10)))]);
2327 assert!(outputs.is_empty());
2328 }
2329
2330 #[mz_ore::test(tokio::test)]
2331 async fn final_output_batch_with_empty_upper() {
2332 let mut inputs = VecDeque::from([input(20, Some(30))]);
2333 let mut outputs = VecDeque::from([(span(20, None), ())]);
2334
2335 let calls = run(
2339 frontier(Some(30)),
2340 &mut inputs,
2341 frontier(None),
2342 &mut outputs,
2343 )
2344 .await;
2345 assert_eq!(calls, vec![Call::Write(span(20, Some(30)), span(20, None))]);
2346 assert!(inputs.is_empty());
2347 assert_eq!(outputs.len(), 1);
2348
2349 let calls = run(frontier(None), &mut inputs, frontier(None), &mut outputs).await;
2350 assert_eq!(calls, vec![Call::Close(span(20, None))]);
2351 assert!(outputs.is_empty());
2352 }
2353 }
2354}
2355
2356fn commit_to_iceberg<'scope>(
2360 name: String,
2361 sink_id: GlobalId,
2362 sink_version: u64,
2363 batch_input: StreamVec<'scope, Timestamp, BoundedDataFile>,
2364 batch_desc_input: StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
2365 table_ready_stream: StreamVec<'scope, Timestamp, Infallible>,
2366 write_frontier: Rc<RefCell<Antichain<Timestamp>>>,
2367 connection: IcebergSinkConnection,
2368 storage_configuration: StorageConfiguration,
2369 write_handle: impl Future<
2370 Output = anyhow::Result<WriteHandle<SourceData, (), Timestamp, StorageDiff>>,
2371 > + 'static,
2372 metrics: Arc<IcebergSinkMetrics>,
2373 statistics: SinkStatistics,
2374) -> (
2375 StreamVec<'scope, Timestamp, HealthStatusMessage>,
2376 PressOnDropButton,
2377) {
2378 let scope = batch_input.scope();
2379 let mut builder = OperatorBuilder::new(name, scope.clone());
2380
2381 let hashed_id = sink_id.hashed();
2382 let is_active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
2383 let name_for_logging = format!("{sink_id}-commit-to-iceberg");
2384
2385 let mut input = builder.new_disconnected_input(batch_input, Exchange::new(move |_| hashed_id));
2386 let mut batch_desc_input =
2387 builder.new_disconnected_input(batch_desc_input, Exchange::new(move |_| hashed_id));
2388 let mut table_ready_input = builder.new_disconnected_input(table_ready_stream, Pipeline);
2389
2390 let (button, errors) = builder.build_fallible(move |_caps| {
2391 Box::pin(async move {
2392 if !is_active_worker {
2393 write_frontier.borrow_mut().clear();
2394 return Ok(());
2395 }
2396
2397 let catalog = connection
2398 .catalog_connection
2399 .connect(&storage_configuration, InTask::Yes)
2400 .await
2401 .with_context(|| {
2402 format!(
2403 "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
2404 connection.catalog_connection.uri, connection.namespace, connection.table
2405 )
2406 })?;
2407
2408 let mut write_handle = write_handle.await?;
2409
2410 let namespace_ident = NamespaceIdent::new(connection.namespace.clone());
2411 let table_ident = TableIdent::new(namespace_ident, connection.table.clone());
2412 while let Some(_) = table_ready_input.next().await {
2413 }
2415 let mut table = catalog.load_table(&table_ident).await.with_context(|| {
2416 format!(
2417 "Failed to load Iceberg table '{}.{}' in commit_to_iceberg operator",
2418 connection.namespace, connection.table
2419 )
2420 })?;
2421
2422 #[allow(clippy::disallowed_types)]
2423 let mut batch_descriptions: std::collections::HashMap<
2424 (Antichain<Timestamp>, Antichain<Timestamp>),
2425 BoundedDataFileSet,
2426 > = std::collections::HashMap::new();
2427
2428 let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
2429 let mut input_frontier = Antichain::from_elem(Timestamp::minimum());
2430
2431 while !(batch_description_frontier.is_empty() && input_frontier.is_empty()) {
2432 tokio::select! {
2433 _ = batch_desc_input.ready() => {},
2434 _ = input.ready() => {}
2435 }
2436
2437 while let Some(event) = batch_desc_input.next_sync() {
2438 match event {
2439 Event::Data(_cap, data) => {
2440 for batch_desc in data {
2441 let prev = batch_descriptions
2442 .insert(batch_desc, BoundedDataFileSet { data_files: vec![] });
2443 if let Some(prev) = prev {
2444 anyhow::bail!(
2445 "Duplicate batch description received \
2446 in commit operator: {:?}",
2447 prev
2448 );
2449 }
2450 }
2451 }
2452 Event::Progress(frontier) => {
2453 batch_description_frontier = frontier;
2454 }
2455 }
2456 }
2457
2458 let ready_events = std::iter::from_fn(|| input.next_sync()).collect_vec();
2459 for event in ready_events {
2460 match event {
2461 Event::Data(_cap, data) => {
2462 for bounded_data_file in data {
2463 let entry = batch_descriptions
2464 .entry(bounded_data_file.batch_desc().clone())
2465 .or_default();
2466 entry.data_files.push(bounded_data_file);
2467 }
2468 }
2469 Event::Progress(frontier) => {
2470 input_frontier = frontier;
2471 }
2472 }
2473 }
2474
2475 let mut done_batches: Vec<_> = batch_descriptions
2481 .keys()
2482 .filter(|(lower, _upper)| PartialOrder::less_than(lower, &input_frontier))
2483 .cloned()
2484 .collect();
2485
2486 done_batches.sort_by(|a, b| {
2488 if PartialOrder::less_than(a, b) {
2489 Ordering::Less
2490 } else if PartialOrder::less_than(b, a) {
2491 Ordering::Greater
2492 } else {
2493 Ordering::Equal
2494 }
2495 });
2496
2497 for batch in done_batches {
2498 let file_set = batch_descriptions.remove(&batch).unwrap();
2499
2500 let mut data_files = vec![];
2501 let mut delete_files = vec![];
2502 let mut total_messages: u64 = 0;
2504 let mut total_bytes: u64 = 0;
2505 for file in file_set.data_files {
2506 total_messages += file.data_file().record_count();
2507 total_bytes += file.data_file().file_size_in_bytes();
2508 match file.data_file().content_type() {
2509 iceberg::spec::DataContentType::Data => {
2510 data_files.push(file.into_data_file());
2511 }
2512 iceberg::spec::DataContentType::PositionDeletes
2513 | iceberg::spec::DataContentType::EqualityDeletes => {
2514 delete_files.push(file.into_data_file());
2515 }
2516 }
2517 }
2518
2519 debug!(
2520 ?sink_id,
2521 %name_for_logging,
2522 lower = %batch.0.pretty(),
2523 upper = %batch.1.pretty(),
2524 data_files = data_files.len(),
2525 delete_files = delete_files.len(),
2526 total_messages,
2527 total_bytes,
2528 "iceberg commit applying batch"
2529 );
2530
2531 let instant = Instant::now();
2532
2533 let frontier = batch.1.clone();
2534 let frontier_json = serde_json::to_string(&frontier.elements())
2535 .context("Failed to serialize frontier to JSON")?;
2536 let snapshot_properties = vec![
2537 ("mz-sink-id".to_string(), sink_id.to_string()),
2538 ("mz-frontier".to_string(), frontier_json),
2539 ("mz-sink-version".to_string(), sink_version.to_string()),
2540 ];
2541
2542 let (table_state, commit_result) = Retry::default()
2543 .max_tries(5)
2544 .retry_async_with_state(table, |_, table| {
2545 let snapshot_properties = snapshot_properties.clone();
2546 let data_files = data_files.clone();
2547 let delete_files = delete_files.clone();
2548 let metrics = Arc::clone(&metrics);
2549 let catalog = Arc::clone(&catalog);
2550 let conn_namespace = connection.namespace.clone();
2551 let conn_table = connection.table.clone();
2552 let frontier = frontier.clone();
2553 let batch_lower = batch.0.clone();
2554 let batch_upper = batch.1.clone();
2555 async move {
2556 try_commit_batch(
2557 table,
2558 snapshot_properties,
2559 data_files,
2560 delete_files,
2561 catalog.as_ref(),
2562 &conn_namespace,
2563 &conn_table,
2564 sink_version,
2565 &frontier,
2566 &batch_lower,
2567 &batch_upper,
2568 &metrics,
2569 )
2570 .await
2571 }
2572 })
2573 .await;
2574 let commit_result = commit_result.with_context(|| {
2575 format!(
2576 "failed to commit batch to Iceberg table '{}.{}'",
2577 connection.namespace, connection.table
2578 )
2579 });
2580 table = table_state;
2581 let duration = instant.elapsed();
2582 metrics
2583 .commit_duration_seconds
2584 .observe(duration.as_secs_f64());
2585 commit_result?;
2586
2587 debug!(
2588 ?sink_id,
2589 %name_for_logging,
2590 lower = %batch.0.pretty(),
2591 upper = %batch.1.pretty(),
2592 total_messages,
2593 total_bytes,
2594 ?duration,
2595 "iceberg commit applied batch"
2596 );
2597
2598 metrics.snapshots_committed.inc();
2599 statistics.inc_messages_committed_by(total_messages);
2600 statistics.inc_bytes_committed_by(total_bytes);
2601
2602 let mut expect_upper = write_handle.shared_upper();
2603 loop {
2604 if PartialOrder::less_equal(&frontier, &expect_upper) {
2605 break;
2607 }
2608
2609 const EMPTY: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
2610 match write_handle
2611 .compare_and_append(EMPTY, expect_upper, frontier.clone())
2612 .await
2613 .expect("valid usage")
2614 {
2615 Ok(()) => break,
2616 Err(mismatch) => {
2617 expect_upper = mismatch.current;
2618 }
2619 }
2620 }
2621 write_frontier.borrow_mut().clone_from(&frontier);
2622 }
2623 }
2624
2625 Ok(())
2626 })
2627 });
2628
2629 let statuses = errors.map(|error| HealthStatusMessage {
2630 id: None,
2631 update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
2632 namespace: StatusNamespace::Iceberg,
2633 });
2634
2635 (statuses, button.press_on_drop())
2636}
2637
2638impl<'scope> SinkRender<'scope> for IcebergSinkConnection {
2639 fn get_key_indices(&self) -> Option<&[usize]> {
2640 self.key_desc_and_indices
2641 .as_ref()
2642 .map(|(_, indices)| indices.as_slice())
2643 }
2644
2645 fn get_relation_key_indices(&self) -> Option<&[usize]> {
2646 self.relation_key_indices.as_deref()
2647 }
2648
2649 fn render_sink(
2650 &self,
2651 storage_state: &mut StorageState,
2652 sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
2653 sink_id: GlobalId,
2654 batches: SinkBatchStream<'scope>,
2655 key_is_synthetic: bool,
2656 _err_collection: VecCollection<'scope, Timestamp, DataflowError, Diff>,
2657 ) -> (
2658 StreamVec<'scope, Timestamp, HealthStatusMessage>,
2659 Vec<PressOnDropButton>,
2660 ) {
2661 let scope = batches.scope();
2662
2663 let write_handle = {
2664 let persist = Arc::clone(&storage_state.persist_clients);
2665 let shard_meta = sink.to_storage_metadata.clone();
2666 async move {
2667 let client = persist.open(shard_meta.persist_location).await?;
2668 let handle = client
2669 .open_writer(
2670 shard_meta.data_shard,
2671 Arc::new(shard_meta.relation_desc),
2672 Arc::new(UnitSchema),
2673 Diagnostics::from_purpose("sink handle"),
2674 )
2675 .await?;
2676 Ok(handle)
2677 }
2678 };
2679
2680 let write_frontier = Rc::new(RefCell::new(Antichain::from_elem(Timestamp::minimum())));
2681 storage_state
2682 .sink_write_frontiers
2683 .insert(sink_id, Rc::clone(&write_frontier));
2684
2685 let (arrow_schema_with_ids, iceberg_schema) =
2686 match (|| -> Result<(ArrowSchema, Arc<Schema>), anyhow::Error> {
2687 let (arrow_schema_with_ids, iceberg_schema) =
2688 relation_desc_to_iceberg_schema(&sink.from_desc)?;
2689
2690 Ok(if sink.envelope == SinkEnvelope::Append {
2691 let extended_arrow = build_schema_with_append_columns(&arrow_schema_with_ids);
2696 let extended_iceberg = Arc::new(
2697 arrow_schema_to_schema(&extended_arrow)
2698 .context("Failed to build Iceberg schema with append columns")?,
2699 );
2700 (extended_arrow, extended_iceberg)
2701 } else {
2702 (arrow_schema_with_ids, iceberg_schema)
2703 })
2704 })() {
2705 Ok(schemas) => schemas,
2706 Err(err) => {
2707 let error_stream = std::iter::once(HealthStatusMessage {
2708 id: None,
2709 update: HealthStatusUpdate::halting(
2710 format!("{}", err.display_with_causes()),
2711 None,
2712 ),
2713 namespace: StatusNamespace::Iceberg,
2714 })
2715 .to_stream(scope);
2716 return (error_stream, vec![]);
2717 }
2718 };
2719
2720 let metrics = Arc::new(
2721 storage_state
2722 .metrics
2723 .get_iceberg_sink_metrics(sink_id, scope.index()),
2724 );
2725
2726 let statistics = storage_state
2727 .aggregated_statistics
2728 .get_sink(&sink_id)
2729 .expect("statistics initialized")
2730 .clone();
2731
2732 let connection_for_minter = self.clone();
2733 let (batch_descriptions, table_ready, mint_status, mint_button) = mint_batch_descriptions(
2734 format!("{sink_id}-iceberg-mint"),
2735 sink_id,
2736 batches.clone(),
2737 sink,
2738 connection_for_minter,
2739 storage_state.storage_configuration.clone(),
2740 Arc::clone(&iceberg_schema),
2741 );
2742
2743 let connection_for_writer = self.clone();
2744 let (datafiles, write_status, write_button) = match sink.envelope {
2745 SinkEnvelope::Upsert => write_data_files::<UpsertEnvelopeHandler>(
2746 format!("{sink_id}-write-data-files"),
2747 batches,
2748 batch_descriptions.clone(),
2749 table_ready.clone(),
2750 sink_id,
2751 sink.from,
2752 key_is_synthetic,
2753 sink.as_of.clone(),
2754 connection_for_writer,
2755 storage_state.storage_configuration.clone(),
2756 Arc::new(arrow_schema_with_ids.clone()),
2757 Arc::clone(&metrics),
2758 statistics.clone(),
2759 ),
2760 SinkEnvelope::Append => write_data_files::<AppendEnvelopeHandler>(
2761 format!("{sink_id}-write-data-files"),
2762 batches,
2763 batch_descriptions.clone(),
2764 table_ready.clone(),
2765 sink_id,
2766 sink.from,
2767 key_is_synthetic,
2768 sink.as_of.clone(),
2769 connection_for_writer,
2770 storage_state.storage_configuration.clone(),
2771 Arc::new(arrow_schema_with_ids.clone()),
2772 Arc::clone(&metrics),
2773 statistics.clone(),
2774 ),
2775 SinkEnvelope::Debezium => {
2776 unreachable!("Iceberg sink only supports Upsert and Append envelopes")
2777 }
2778 };
2779
2780 let connection_for_committer = self.clone();
2781 let (commit_status, commit_button) = commit_to_iceberg(
2782 format!("{sink_id}-commit-to-iceberg"),
2783 sink_id,
2784 sink.version,
2785 datafiles,
2786 batch_descriptions,
2787 table_ready,
2788 Rc::clone(&write_frontier),
2789 connection_for_committer,
2790 storage_state.storage_configuration.clone(),
2791 write_handle,
2792 Arc::clone(&metrics),
2793 statistics,
2794 );
2795
2796 let running_status = Some(HealthStatusMessage {
2797 id: None,
2798 update: HealthStatusUpdate::running(),
2799 namespace: StatusNamespace::Iceberg,
2800 })
2801 .to_stream(scope);
2802
2803 let statuses =
2804 scope.concatenate([running_status, mint_status, write_status, commit_status]);
2805
2806 (statuses, vec![mint_button, write_button, commit_button])
2807 }
2808}