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 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 let table_creation = TableCreation::builder()
956 .name(table_name.clone())
957 .schema(schema.clone())
958 .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 Err(err).context("Failed to load Iceberg table")
975 }
976 }
977 }
978}
979
980fn 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 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
1020fn 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
1044fn 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
1076fn 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#[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
1105fn 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 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 let overcompacted =
1201 *resume_upper != [Timestamp::minimum()] &&
1203 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 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 let mut minted_batches = VecDeque::new();
1229
1230 let catchup_start = if *resume_upper == [Timestamp::minimum()] {
1233 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 batch_upper
1243 } else {
1244 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 if catchup_start.is_empty() {
1269 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 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 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 return Ok(());
1349 }
1350 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#[derive(Clone, Debug, Serialize, Deserialize)]
1409struct AvroDataFile {
1410 pub data_file: Vec<u8>,
1411 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#[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#[derive(Clone, Debug, Default)]
1488struct BoundedDataFileSet {
1489 pub data_files: Vec<BoundedDataFile>,
1490}
1491
1492fn data_file_location(configured_path: Option<&str>, location: &str) -> String {
1502 if let Some(path) = configured_path {
1509 return path.trim_end_matches('/').to_string();
1510 }
1511
1512 let corrected_location = match location.rsplit_once("/metadata/") {
1516 Some((a, b)) if b.ends_with(".metadata.json") => a,
1517 _ => location,
1518 };
1519 format!("{}/data", corrected_location.trim_end_matches('/'))
1522}
1523
1524fn 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 }
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 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 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 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(¤t_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 let mut stashed_rows: VecDeque<ArcBatch<OrdValBatch<_>>> = VecDeque::new();
1651
1652 let mut in_flight_batches: VecDeque<(
1656 (Antichain<Timestamp>, Antichain<Timestamp>),
1657 Box<dyn IcebergWriter>,
1658 )> = VecDeque::new();
1659
1660 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 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 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 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 !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 let mut staged_messages_since_flush: u64 = 0;
1762
1763 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 if let Some(warner) = pk_warner.as_mut() {
1798 warner.flush();
1799 }
1800 Ok(())
1801 };
1802
1803 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 capset.downgrade(batch_desc.1.clone());
1842 Ok(())
1843 };
1844
1845 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
1875type BatchDescription = (Antichain<Timestamp>, Antichain<Timestamp>);
1877
1878async 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 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 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 break;
1933 };
1934
1935 let Some(rows) = input_batches.front() else {
1936 break;
1938 };
1939
1940 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 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 input_batches.pop_front();
1959 }
1960
1961 }
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 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 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 assert_eq!(
2003 manifest_uri(Some("s3://bucket/tbl/data"), "s3://bucket/tbl"),
2004 "s3://bucket/tbl/data/part-00000.parquet"
2005 );
2006
2007 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 assert_eq!(
2023 manifest_uri(None, "s3://bucket/tbl"),
2024 "s3://bucket/tbl/data/part-00000.parquet"
2025 );
2026
2027 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 assert_eq!(
2039 data_file_location(None, "s3://bucket/tbl/metadata/00001-abc.metadata.json"),
2040 "s3://bucket/tbl/data"
2041 );
2042
2043 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 let result = iceberg_type_overrides(&SqlScalarType::UInt16);
2054 assert_eq!(result.unwrap().0, DataType::Int32);
2055
2056 let result = iceberg_type_overrides(&SqlScalarType::UInt32);
2058 assert_eq!(result.unwrap().0, DataType::Int64);
2059
2060 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 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 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 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 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 let result = iceberg_type_overrides(&SqlScalarType::Interval);
2114 assert_eq!(result.unwrap().0, DataType::LargeUtf8);
2115
2116 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 assert_eq!(arrow_schema.field(1).data_type(), &DataType::LargeUtf8);
2127
2128 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 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 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 #[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 assert_eq!(entry_fields[0].name(), "key");
2255 assert_eq!(entry_fields[1].name(), "value");
2256 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 fn frontier(t: Option<u64>) -> Antichain<Timestamp> {
2274 t.map_or_else(Antichain::new, |t| Antichain::from_elem(Timestamp::new(t)))
2275 }
2276
2277 fn span(lower: u64, upper: Option<u64>) -> BatchDescription {
2279 (frontier(Some(lower)), frontier(upper))
2280 }
2281
2282 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 Write(BatchDescription, BatchDescription),
2293 Close(BatchDescription),
2294 }
2295
2296 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 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 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 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 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 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 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
2488fn 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 }
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 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 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 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 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 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}