1use std::cmp::Ordering;
86use std::collections::{BTreeSet, VecDeque};
87use std::convert::Infallible;
88use std::future::Future;
89use std::time::Instant;
90use std::{cell::RefCell, rc::Rc, sync::Arc};
91
92use anyhow::{Context, anyhow};
93use arrow::array::{ArrayRef, Int32Array, Int64Array, RecordBatch};
94use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
95use differential_dataflow::trace::Span;
96use differential_dataflow::trace::implementations::ord_neu::OrdValBatch;
97use differential_dataflow::trace::implementations::spine_fueled::SpineBatch;
98use differential_dataflow::trace::implementations::{BatchContainer, Layout};
99use differential_dataflow::{Hashable, VecCollection};
100use futures::StreamExt;
101use iceberg::ErrorKind;
102use iceberg::arrow::{
103 arrow_schema_to_schema, arrow_schema_to_schema_auto_assign_ids, schema_to_arrow_schema,
104};
105use iceberg::spec::{
106 DataFile, FormatVersion, NestedField, NestedFieldRef, PrimitiveType, Snapshot, StructType,
107 Type, read_data_files_from_avro, write_data_files_to_avro,
108};
109use iceberg::spec::{Schema, SchemaRef};
110use iceberg::table::Table;
111use iceberg::transaction::{RowDeltaAction, TransactionAction};
112use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
113use iceberg::writer::base_writer::equality_delete_writer::{
114 EqualityDeleteFileWriterBuilder, EqualityDeleteWriterConfig,
115};
116use iceberg::writer::base_writer::position_delete_writer::{
117 PositionDeleteFileWriterBuilder, PositionDeleteWriterConfig,
118};
119use iceberg::writer::combined_writer::delta_writer::DeltaWriterBuilder;
120use iceberg::writer::file_writer::ParquetWriterBuilder;
121use iceberg::writer::file_writer::location_generator::{
122 DefaultFileNameGenerator, DefaultLocationGenerator,
123};
124use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
125use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
126use iceberg::{Catalog, NamespaceIdent, TableCommit, TableCreation, TableIdent};
127use itertools::Itertools;
128use mz_arrow_util::builder::{ARROW_EXTENSION_NAME_KEY, ArrowBuilder};
129use mz_interchange::avro::DiffPair;
130use mz_interchange::envelopes::for_each_diff_pair_async;
131use mz_ore::cast::CastFrom;
132use mz_ore::error::ErrorExt;
133use mz_ore::future::InTask;
134use mz_ore::result::ResultExt;
135use mz_ore::retry::{Retry, RetryResult};
136use mz_persist_client::Diagnostics;
137use mz_persist_client::write::WriteHandle;
138use mz_persist_types::codec_impls::UnitSchema;
139use mz_repr::{Diff, GlobalId, Row, Timestamp};
140use mz_row_spine::ArcBatch;
141use mz_storage_types::StorageDiff;
142use mz_storage_types::configuration::StorageConfiguration;
143use mz_storage_types::controller::CollectionMetadata;
144use mz_storage_types::errors::DataflowError;
145use mz_storage_types::sinks::{
146 IcebergSinkConnection, SinkEnvelope, StorageSinkDesc, iceberg_type_overrides,
147};
148use mz_storage_types::sources::SourceData;
149use mz_timely_util::antichain::AntichainExt;
150use mz_timely_util::builder_async::{Event, OperatorBuilder, PressOnDropButton};
151use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
152use parquet::file::properties::WriterProperties;
153use serde::{Deserialize, Serialize};
154use timely::PartialOrder;
155use timely::container::CapacityContainerBuilder;
156use timely::dataflow::StreamVec;
157use timely::dataflow::channels::pact::{Exchange, Pipeline};
158use timely::dataflow::operators::vec::{Broadcast, Map, ToStream};
159use timely::dataflow::operators::{CapabilitySet, Concatenate};
160use timely::progress::{Antichain, Timestamp as _};
161use tracing::{debug, info};
162
163use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace};
164use crate::metrics::sink::iceberg::IcebergSinkMetrics;
165use crate::render::sinks::{PkViolationWarner, SinkBatchStream, SinkRender};
166use crate::statistics::SinkStatistics;
167use crate::storage_state::StorageState;
168
169const DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY: usize = 1024;
172const DEFAULT_ARRAY_BUILDER_DATA_CAPACITY: usize = 1024;
176
177const PARQUET_FILE_PREFIX: &str = "mz_data";
179const INITIAL_DESCRIPTIONS_TO_MINT: u64 = 3;
182
183struct WriterContext {
186 arrow_schema: Arc<ArrowSchema>,
188 current_schema: Arc<Schema>,
190 file_io: iceberg::io::FileIO,
192 location_generator: DefaultLocationGenerator,
194 file_name_generator: DefaultFileNameGenerator,
196 writer_properties: WriterProperties,
197}
198
199trait EnvelopeHandler: Send {
201 fn new(
203 ctx: WriterContext,
204 connection: &IcebergSinkConnection,
205 materialize_arrow_schema: &Arc<ArrowSchema>,
206 ) -> anyhow::Result<Self>
207 where
208 Self: Sized;
209
210 async fn create_writer(&self, is_snapshot: bool) -> anyhow::Result<Box<dyn IcebergWriter>>;
216
217 fn row_to_batch(&self, diff_pair: DiffPair<Row>, ts: Timestamp) -> anyhow::Result<RecordBatch>;
218}
219
220struct UpsertEnvelopeHandler {
221 ctx: WriterContext,
222 equality_ids: Vec<i32>,
224 pos_schema: Arc<Schema>,
226 eq_schema: Arc<Schema>,
228 eq_config: EqualityDeleteWriterConfig,
230 schema_with_op: Arc<ArrowSchema>,
234}
235
236impl EnvelopeHandler for UpsertEnvelopeHandler {
237 fn new(
238 ctx: WriterContext,
239 connection: &IcebergSinkConnection,
240 materialize_arrow_schema: &Arc<ArrowSchema>,
241 ) -> anyhow::Result<Self> {
242 let Some((_, equality_indices)) = &connection.key_desc_and_indices else {
243 return Err(anyhow::anyhow!(
244 "Iceberg sink requires key columns for equality deletes"
245 ));
246 };
247
248 let equality_ids = equality_ids_for_indices(
249 ctx.current_schema.as_ref(),
250 materialize_arrow_schema.as_ref(),
251 equality_indices,
252 )?;
253
254 let pos_arrow_schema = PositionDeleteWriterConfig::arrow_schema();
255 let pos_schema = Arc::new(
256 arrow_schema_to_schema(&pos_arrow_schema)
257 .context("Failed to convert position delete Arrow schema to Iceberg schema")?,
258 );
259
260 let eq_config =
261 EqualityDeleteWriterConfig::new(equality_ids.clone(), Arc::clone(&ctx.current_schema))
262 .context("Failed to create EqualityDeleteWriterConfig")?;
263 let eq_schema = Arc::new(
264 arrow_schema_to_schema(eq_config.projected_arrow_schema_ref())
265 .context("Failed to convert equality delete Arrow schema to Iceberg schema")?,
266 );
267
268 let schema_with_op = Arc::new(build_schema_with_op_column(&ctx.arrow_schema));
269
270 Ok(Self {
271 ctx,
272 equality_ids,
273 pos_schema,
274 eq_schema,
275 eq_config,
276 schema_with_op,
277 })
278 }
279
280 async fn create_writer(&self, is_snapshot: bool) -> anyhow::Result<Box<dyn IcebergWriter>> {
281 let data_parquet_writer = ParquetWriterBuilder::new(
282 self.ctx.writer_properties.clone(),
283 Arc::clone(&self.ctx.current_schema),
284 )
285 .with_arrow_schema(Arc::clone(&self.ctx.arrow_schema))
286 .context("Arrow schema validation failed")?;
287 let data_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
288 data_parquet_writer,
289 Arc::clone(&self.ctx.current_schema),
290 self.ctx.file_io.clone(),
291 self.ctx.location_generator.clone(),
292 self.ctx.file_name_generator.clone(),
293 );
294 let data_writer_builder = DataFileWriterBuilder::new(data_rolling_writer);
295
296 let pos_config = PositionDeleteWriterConfig::new(None, 0, None);
297 let pos_parquet_writer = ParquetWriterBuilder::new(
298 self.ctx.writer_properties.clone(),
299 Arc::clone(&self.pos_schema),
300 );
301 let pos_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
302 pos_parquet_writer,
303 Arc::clone(&self.ctx.current_schema),
304 self.ctx.file_io.clone(),
305 self.ctx.location_generator.clone(),
306 self.ctx.file_name_generator.clone(),
307 );
308 let pos_delete_writer_builder =
309 PositionDeleteFileWriterBuilder::new(pos_rolling_writer, pos_config);
310
311 let eq_parquet_writer = ParquetWriterBuilder::new(
312 self.ctx.writer_properties.clone(),
313 Arc::clone(&self.eq_schema),
314 );
315 let eq_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
316 eq_parquet_writer,
317 Arc::clone(&self.ctx.current_schema),
318 self.ctx.file_io.clone(),
319 self.ctx.location_generator.clone(),
320 self.ctx.file_name_generator.clone(),
321 );
322 let eq_delete_writer_builder =
323 EqualityDeleteFileWriterBuilder::new(eq_rolling_writer, self.eq_config.clone());
324
325 let mut builder = DeltaWriterBuilder::new(
326 data_writer_builder,
327 pos_delete_writer_builder,
328 eq_delete_writer_builder,
329 self.equality_ids.clone(),
330 );
331
332 builder = if is_snapshot {
333 builder.with_max_seen_rows(0)
335 } else {
336 builder.with_max_seen_rows(usize::MAX)
348 };
349
350 Ok(Box::new(
351 builder
352 .build(None)
353 .await
354 .context("Failed to create DeltaWriter")?,
355 ))
356 }
357
358 fn row_to_batch(
361 &self,
362 diff_pair: DiffPair<Row>,
363 _ts: Timestamp,
364 ) -> anyhow::Result<RecordBatch> {
365 let mut builder = ArrowBuilder::new_with_schema(
366 Arc::clone(&self.ctx.arrow_schema),
367 DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY,
368 DEFAULT_ARRAY_BUILDER_DATA_CAPACITY,
369 )
370 .context("Failed to create builder")?;
371
372 let mut op_values = Vec::new();
373
374 if let Some(before) = diff_pair.before {
375 builder
376 .add_row(&before)
377 .context("Failed to add delete row to builder")?;
378 op_values.push(-1i32);
379 }
380 if let Some(after) = diff_pair.after {
381 builder
382 .add_row(&after)
383 .context("Failed to add insert row to builder")?;
384 op_values.push(1i32);
385 }
386
387 let batch = builder
388 .to_record_batch()
389 .context("Failed to create record batch")?;
390
391 let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
392 columns.push(Arc::new(Int32Array::from(op_values)));
393
394 RecordBatch::try_new(Arc::clone(&self.schema_with_op), columns)
395 .context("Failed to create batch with op column")
396 }
397}
398
399struct AppendEnvelopeHandler {
400 ctx: WriterContext,
401 user_schema_for_append: Arc<ArrowSchema>,
404}
405
406impl EnvelopeHandler for AppendEnvelopeHandler {
407 fn new(
408 ctx: WriterContext,
409 _connection: &IcebergSinkConnection,
410 _materialize_arrow_schema: &Arc<ArrowSchema>,
411 ) -> anyhow::Result<Self> {
412 let n = ctx.arrow_schema.fields().len().saturating_sub(2);
415 let user_schema_for_append =
416 Arc::new(ArrowSchema::new(ctx.arrow_schema.fields()[..n].to_vec()));
417
418 Ok(Self {
419 ctx,
420 user_schema_for_append,
421 })
422 }
423
424 async fn create_writer(&self, _is_snapshot: bool) -> anyhow::Result<Box<dyn IcebergWriter>> {
425 let data_parquet_writer = ParquetWriterBuilder::new(
426 self.ctx.writer_properties.clone(),
427 Arc::clone(&self.ctx.current_schema),
428 )
429 .with_arrow_schema(Arc::clone(&self.ctx.arrow_schema))
430 .context("Arrow schema validation failed")?;
431 let data_rolling_writer = RollingFileWriterBuilder::new_with_default_file_size(
432 data_parquet_writer,
433 Arc::clone(&self.ctx.current_schema),
434 self.ctx.file_io.clone(),
435 self.ctx.location_generator.clone(),
436 self.ctx.file_name_generator.clone(),
437 );
438 Ok(Box::new(
439 DataFileWriterBuilder::new(data_rolling_writer)
440 .build(None)
441 .await
442 .context("Failed to create DataFileWriter")?,
443 ))
444 }
445
446 fn row_to_batch(&self, diff_pair: DiffPair<Row>, ts: Timestamp) -> anyhow::Result<RecordBatch> {
449 let mut builder = ArrowBuilder::new_with_schema(
450 Arc::clone(&self.user_schema_for_append),
451 DEFAULT_ARRAY_BUILDER_ITEM_CAPACITY,
452 DEFAULT_ARRAY_BUILDER_DATA_CAPACITY,
453 )
454 .context("Failed to create builder")?;
455
456 let mut diff_values: Vec<i32> = Vec::new();
457 let ts_i64 = i64::try_from(u64::from(ts)).unwrap_or(i64::MAX);
458
459 if let Some(before) = diff_pair.before {
460 builder
461 .add_row(&before)
462 .context("Failed to add before row to builder")?;
463 diff_values.push(-1i32);
464 }
465 if let Some(after) = diff_pair.after {
466 builder
467 .add_row(&after)
468 .context("Failed to add after row to builder")?;
469 diff_values.push(1i32);
470 }
471
472 let n = diff_values.len();
473 let batch = builder
474 .to_record_batch()
475 .context("Failed to create record batch")?;
476
477 let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
478 columns.push(Arc::new(Int32Array::from(diff_values)));
479 columns.push(Arc::new(Int64Array::from(vec![ts_i64; n])));
480
481 RecordBatch::try_new(Arc::clone(&self.ctx.arrow_schema), columns)
482 .context("Failed to create append record batch")
483 }
484}
485
486fn add_field_ids_to_arrow_schema(schema: ArrowSchema) -> ArrowSchema {
491 let mut next_field_id = 1i32;
492 let fields: Vec<Field> = schema
493 .fields()
494 .iter()
495 .map(|field| add_field_ids_recursive(field, &mut next_field_id))
496 .collect();
497 ArrowSchema::new(fields).with_metadata(schema.metadata().clone())
498}
499
500fn add_field_ids_recursive(field: &Field, next_id: &mut i32) -> Field {
502 let current_id = *next_id;
503 *next_id += 1;
504
505 let mut metadata = field.metadata().clone();
506 metadata.insert(
507 PARQUET_FIELD_ID_META_KEY.to_string(),
508 current_id.to_string(),
509 );
510
511 let new_data_type = add_field_ids_to_datatype(field.data_type(), next_id);
512
513 Field::new(field.name(), new_data_type, field.is_nullable()).with_metadata(metadata)
514}
515
516fn add_field_ids_to_datatype(data_type: &DataType, next_id: &mut i32) -> DataType {
518 match data_type {
519 DataType::Struct(fields) => {
520 let new_fields: Vec<Field> = fields
521 .iter()
522 .map(|f| add_field_ids_recursive(f, next_id))
523 .collect();
524 DataType::Struct(new_fields.into())
525 }
526 DataType::List(element_field) => {
527 let new_element = add_field_ids_recursive(element_field, next_id);
528 DataType::List(Arc::new(new_element))
529 }
530 DataType::LargeList(element_field) => {
531 let new_element = add_field_ids_recursive(element_field, next_id);
532 DataType::LargeList(Arc::new(new_element))
533 }
534 DataType::Map(entries_field, sorted) => {
535 let new_entries = add_field_ids_recursive(entries_field, next_id);
536 DataType::Map(Arc::new(new_entries), *sorted)
537 }
538 _ => data_type.clone(),
539 }
540}
541
542fn merge_materialize_metadata_into_iceberg_schema(
547 materialize_arrow_schema: &ArrowSchema,
548 iceberg_schema: &Schema,
549) -> anyhow::Result<ArrowSchema> {
550 let iceberg_arrow_schema = schema_to_arrow_schema(iceberg_schema)
552 .context("Failed to convert Iceberg schema to Arrow schema")?;
553
554 let fields: Vec<Field> = iceberg_arrow_schema
556 .fields()
557 .iter()
558 .map(|iceberg_field| {
559 let mz_field = materialize_arrow_schema
561 .field_with_name(iceberg_field.name())
562 .with_context(|| {
563 format!(
564 "Field '{}' not found in Materialize schema",
565 iceberg_field.name()
566 )
567 })?;
568
569 merge_field_metadata_recursive(iceberg_field, Some(mz_field))
570 })
571 .collect::<anyhow::Result<Vec<_>>>()?;
572
573 Ok(ArrowSchema::new(fields).with_metadata(iceberg_arrow_schema.metadata().clone()))
574}
575
576fn merge_field_metadata_recursive(
578 iceberg_field: &Field,
579 mz_field: Option<&Field>,
580) -> anyhow::Result<Field> {
581 let mut metadata = iceberg_field.metadata().clone();
583
584 if let Some(mz_f) = mz_field {
586 if let Some(extension_name) = mz_f.metadata().get(ARROW_EXTENSION_NAME_KEY) {
587 metadata.insert(ARROW_EXTENSION_NAME_KEY.to_string(), extension_name.clone());
588 }
589 }
590
591 let new_data_type = match iceberg_field.data_type() {
593 DataType::Struct(iceberg_fields) => {
594 let mz_struct_fields = match mz_field {
595 Some(f) => match f.data_type() {
596 DataType::Struct(fields) => Some(fields),
597 other => anyhow::bail!(
598 "Type mismatch for field '{}': Iceberg schema has Struct, but Materialize schema has {:?}",
599 iceberg_field.name(),
600 other
601 ),
602 },
603 None => None,
604 };
605
606 let new_fields: Vec<Field> = iceberg_fields
607 .iter()
608 .map(|iceberg_inner| {
609 let mz_inner = mz_struct_fields.and_then(|fields| {
610 fields.iter().find(|f| f.name() == iceberg_inner.name())
611 });
612 merge_field_metadata_recursive(iceberg_inner, mz_inner.map(|f| f.as_ref()))
613 })
614 .collect::<anyhow::Result<Vec<_>>>()?;
615
616 DataType::Struct(new_fields.into())
617 }
618 DataType::List(iceberg_element) => {
619 let mz_element = match mz_field {
620 Some(f) => match f.data_type() {
621 DataType::List(element) => Some(element.as_ref()),
622 other => anyhow::bail!(
623 "Type mismatch for field '{}': Iceberg schema has List, but Materialize schema has {:?}",
624 iceberg_field.name(),
625 other
626 ),
627 },
628 None => None,
629 };
630 let new_element = merge_field_metadata_recursive(iceberg_element, mz_element)?;
631 DataType::List(Arc::new(new_element))
632 }
633 DataType::LargeList(iceberg_element) => {
634 let mz_element = match mz_field {
635 Some(f) => match f.data_type() {
636 DataType::LargeList(element) => Some(element.as_ref()),
637 other => anyhow::bail!(
638 "Type mismatch for field '{}': Iceberg schema has LargeList, but Materialize schema has {:?}",
639 iceberg_field.name(),
640 other
641 ),
642 },
643 None => None,
644 };
645 let new_element = merge_field_metadata_recursive(iceberg_element, mz_element)?;
646 DataType::LargeList(Arc::new(new_element))
647 }
648 DataType::Map(iceberg_entries, sorted) => {
649 let mz_entries = match mz_field {
650 Some(f) => match f.data_type() {
651 DataType::Map(entries, _) => Some(entries.as_ref()),
652 other => anyhow::bail!(
653 "Type mismatch for field '{}': Iceberg schema has Map, but Materialize schema has {:?}",
654 iceberg_field.name(),
655 other
656 ),
657 },
658 None => None,
659 };
660 let new_entries = match mz_entries {
665 Some(mz_entries) => merge_map_entries_metadata(iceberg_entries, mz_entries)?,
666 None => iceberg_entries.as_ref().clone(),
667 };
668 DataType::Map(Arc::new(new_entries), *sorted)
669 }
670 other => other.clone(),
671 };
672
673 Ok(Field::new(
674 iceberg_field.name(),
675 new_data_type,
676 iceberg_field.is_nullable(),
677 )
678 .with_metadata(metadata))
679}
680
681fn merge_map_entries_metadata(
703 iceberg_entries: &Field,
704 mz_entries: &Field,
705) -> anyhow::Result<Field> {
706 let mut metadata = iceberg_entries.metadata().clone();
707 if let Some(extension_name) = mz_entries.metadata().get(ARROW_EXTENSION_NAME_KEY) {
708 metadata.insert(ARROW_EXTENSION_NAME_KEY.to_string(), extension_name.clone());
709 }
710
711 let iceberg_fields = match iceberg_entries.data_type() {
712 DataType::Struct(fields) => fields,
713 other => anyhow::bail!(
714 "Iceberg map entries field '{}' is not a Struct: {:?}",
715 iceberg_entries.name(),
716 other
717 ),
718 };
719 let mz_fields = match mz_entries.data_type() {
720 DataType::Struct(fields) => fields,
721 other => anyhow::bail!(
722 "Materialize map entries field '{}' is not a Struct: {:?}",
723 mz_entries.name(),
724 other
725 ),
726 };
727
728 let new_fields: Vec<Field> = iceberg_fields
729 .iter()
730 .enumerate()
731 .map(|(idx, iceberg_inner)| {
732 let mz_inner = mz_fields.get(idx).map(|f| f.as_ref());
733 merge_field_metadata_recursive(iceberg_inner, mz_inner)
734 })
735 .collect::<anyhow::Result<Vec<_>>>()?;
736
737 Ok(Field::new(
738 iceberg_entries.name(),
739 DataType::Struct(new_fields.into()),
740 iceberg_entries.is_nullable(),
741 )
742 .with_metadata(metadata))
743}
744
745async fn reload_table(
746 catalog: &dyn Catalog,
747 namespace: String,
748 table_name: String,
749 current_table: &Table,
750) -> anyhow::Result<Table> {
751 let namespace_ident = NamespaceIdent::new(namespace.clone());
752 let table_ident = TableIdent::new(namespace_ident, table_name.clone());
753 let current_schema = current_table.metadata().current_schema_id();
754 let current_partition_spec = current_table.metadata().default_partition_spec_id();
755
756 match catalog.load_table(&table_ident).await {
757 Ok(table) => {
758 let reloaded_schema = table.metadata().current_schema_id();
759 let reloaded_partition_spec = table.metadata().default_partition_spec_id();
760 if reloaded_schema != current_schema {
761 return Err(anyhow::anyhow!(
762 "Iceberg table '{}' schema changed during operation but schema evolution isn't supported, expected schema ID {}, got {}",
763 table_name,
764 current_schema,
765 reloaded_schema
766 ));
767 }
768
769 if reloaded_partition_spec != current_partition_spec {
770 return Err(anyhow::anyhow!(
771 "Iceberg table '{}' partition spec changed during operation but partition spec evolution isn't supported, expected partition spec ID {}, got {}",
772 table_name,
773 current_partition_spec,
774 reloaded_partition_spec
775 ));
776 }
777
778 Ok(table)
779 }
780 Err(err) => Err(err).context("Failed to reload Iceberg table"),
781 }
782}
783
784enum CommitError {
786 Local(iceberg::Error),
789 Request(iceberg::Error),
792}
793
794async fn do_commit(
805 table: &Table,
806 catalog: &dyn Catalog,
807 snapshot_properties: Vec<(String, String)>,
808 data_files: Vec<DataFile>,
809 delete_files: Vec<DataFile>,
810) -> Result<Table, CommitError> {
811 let mut action = RowDeltaAction::new()
812 .set_snapshot_properties(snapshot_properties.into_iter().collect())
813 .with_check_duplicate(false);
814
815 if !data_files.is_empty() || !delete_files.is_empty() {
816 action = action
817 .add_data_files(data_files)
818 .add_delete_files(delete_files);
819 }
820
821 let mut action_commit = Arc::new(action)
827 .commit(table)
828 .await
829 .map_err(CommitError::Local)?;
830
831 let table_commit = TableCommit::builder()
837 .ident(table.identifier().clone())
838 .updates(action_commit.take_updates())
839 .requirements(action_commit.take_requirements())
840 .build();
841
842 catalog
843 .update_table(table_commit)
844 .await
845 .map_err(CommitError::Request)
846}
847
848async fn try_commit_batch(
850 table: Table,
851 snapshot_properties: Vec<(String, String)>,
852 data_files: Vec<DataFile>,
853 delete_files: Vec<DataFile>,
854 catalog: &dyn Catalog,
855 conn_namespace: &str,
856 conn_table: &str,
857 sink_id: GlobalId,
858 sink_version: u64,
859 batch_lower: &Antichain<Timestamp>,
860 batch_upper: &Antichain<Timestamp>,
861 metrics: &IcebergSinkMetrics,
862) -> (Table, RetryResult<(), anyhow::Error>) {
863 let table = match reload_table(
867 catalog,
868 conn_namespace.to_string(),
869 conn_table.to_string(),
870 &table,
871 )
872 .await
873 {
874 Ok(table) => table,
875 Err(e) => {
876 return (table, RetryResult::RetryableErr(anyhow!(e)));
878 }
879 };
880
881 let mut snapshots: Vec<_> = table.metadata().snapshots().cloned().collect();
882 let last = match retrieve_upper_from_snapshots(&mut snapshots) {
883 Ok(last) => last,
884 Err(e) => return (table, RetryResult::RetryableErr(anyhow!(e))),
885 };
886 if let Some((last_frontier, last_id, last_version)) = last {
887 if last_id == sink_id && last_version == sink_version && last_frontier == *batch_upper {
889 info!(
892 namespace = %conn_namespace,
893 table = %conn_table,
894 lower = %batch_lower.pretty(),
895 upper = %batch_upper.pretty(),
896 "found iceberg commit from previous attempt, treating as success"
897 );
898 return (table, RetryResult::Ok(()));
899 }
900
901 if last_version > sink_version {
902 return (
904 table,
905 RetryResult::FatalErr(anyhow!(
906 "Iceberg table '{}' has been modified by another writer \
907 with version {}. Current sink version: {}. \
908 Frontiers may be out of sync, aborting to avoid data loss.",
909 conn_table,
910 last_version,
911 sink_version,
912 )),
913 );
914 }
915
916 if PartialOrder::less_equal(batch_upper, &last_frontier)
917 || PartialOrder::less_than(batch_lower, &last_frontier)
918 {
919 return (
921 table,
922 RetryResult::FatalErr(anyhow!(
923 "Iceberg table '{}' has been modified by another writer. \
924 Current frontier: {:?}, last frontier: {:?}.",
925 conn_table,
926 batch_upper,
927 last_frontier,
928 )),
929 );
930 }
931 }
932
933 match do_commit(
934 &table,
935 catalog,
936 snapshot_properties,
937 data_files,
938 delete_files,
939 )
940 .await
941 {
942 Ok(new_table) => (new_table, RetryResult::Ok(())),
943 Err(CommitError::Local(e)) => {
944 metrics.commit_failures.inc();
947 (
948 table,
949 RetryResult::RetryableErr(anyhow!("Failed to build iceberg table commit: {}", e)),
950 )
951 }
952 Err(CommitError::Request(e)) => match e.kind() {
953 ErrorKind::CatalogCommitConflicts => {
954 metrics.commit_conflicts.inc();
957 (table, RetryResult::RetryableErr(anyhow!(e)))
958 }
959 ErrorKind::Unexpected => {
960 metrics.commit_failures.inc();
963 (table, RetryResult::RetryableErr(anyhow!(e)))
964 }
965 _ => {
966 metrics.commit_failures.inc();
968 (table, RetryResult::FatalErr(anyhow!(e)))
969 }
970 },
971 }
972}
973
974fn schema_mismatch(current: &Schema, expected: &Schema) -> Option<String> {
994 struct_mismatch(None, current.as_struct(), expected.as_struct())
995}
996
997fn struct_mismatch(
998 parent: Option<&str>,
999 current: &StructType,
1000 expected: &StructType,
1001) -> Option<String> {
1002 let (current, expected) = (current.fields(), expected.fields());
1003 let names = |fields: &[NestedFieldRef]| -> BTreeSet<String> {
1006 fields.iter().map(|field| field.name.clone()).collect()
1007 };
1008 let (current_names, expected_names) = (names(current), names(expected));
1009 if let Some(missing) = expected
1010 .iter()
1011 .find(|field| !current_names.contains(&field.name))
1012 {
1013 return Some(format!(
1014 "the table is missing `{}`",
1015 field_path(parent, &missing.name)
1016 ));
1017 }
1018 if let Some(extra) = current
1019 .iter()
1020 .find(|field| !expected_names.contains(&field.name))
1021 {
1022 return Some(format!(
1023 "the table has an extra field `{}`",
1024 field_path(parent, &extra.name)
1025 ));
1026 }
1027 std::iter::zip(current, expected).find_map(|(current_field, expected_field)| {
1030 field_mismatch(parent, current_field, expected_field)
1031 })
1032}
1033
1034fn field_mismatch(
1035 parent: Option<&str>,
1036 current: &NestedField,
1037 expected: &NestedField,
1038) -> Option<String> {
1039 let path = field_path(parent, &expected.name);
1040 if current.name != expected.name {
1041 return Some(format!(
1042 "the table has `{}` where `{path}` is expected",
1043 field_path(parent, ¤t.name)
1044 ));
1045 }
1046 if current.required != expected.required {
1047 let nullability = |required| if required { "required" } else { "optional" };
1048 return Some(format!(
1049 "`{path}` is {} in the table but {} in the sink",
1050 nullability(current.required),
1051 nullability(expected.required)
1052 ));
1053 }
1054 match (&*current.field_type, &*expected.field_type) {
1055 (Type::Primitive(PrimitiveType::Fixed(16)), Type::Primitive(PrimitiveType::String)) => None,
1057 (Type::Struct(current_struct), Type::Struct(expected_struct)) => {
1058 struct_mismatch(Some(&path), current_struct, expected_struct)
1059 }
1060 (Type::List(current_list), Type::List(expected_list)) => field_mismatch(
1061 Some(&path),
1062 ¤t_list.element_field,
1063 &expected_list.element_field,
1064 ),
1065 (Type::Map(current_map), Type::Map(expected_map)) => {
1066 let parent = Some(path.as_str());
1067 field_mismatch(parent, ¤t_map.key_field, &expected_map.key_field).or_else(|| {
1068 field_mismatch(parent, ¤t_map.value_field, &expected_map.value_field)
1069 })
1070 }
1071 (current_type, expected_type) if current_type == expected_type => None,
1072 (current_type, expected_type) => Some(format!(
1073 "`{path}` has type {} in the table but {} in the sink",
1074 type_name(current_type),
1075 type_name(expected_type)
1076 )),
1077 }
1078}
1079
1080fn field_path(parent: Option<&str>, name: &str) -> String {
1081 match parent {
1082 Some(parent) => format!("{parent}.{name}"),
1083 None => name.to_string(),
1084 }
1085}
1086
1087fn type_name(ty: &Type) -> String {
1090 match ty {
1091 Type::Primitive(primitive) => primitive.to_string(),
1092 Type::Struct(_) => "struct".to_string(),
1093 Type::List(_) => "list".to_string(),
1094 Type::Map(_) => "map".to_string(),
1095 }
1096}
1097
1098async fn load_or_create_table(
1100 catalog: &dyn Catalog,
1101 namespace: String,
1102 table_name: String,
1103 schema: &Schema,
1104) -> anyhow::Result<iceberg::table::Table> {
1105 let namespace_ident = NamespaceIdent::new(namespace.clone());
1106 let table_ident = TableIdent::new(namespace_ident.clone(), table_name.clone());
1107
1108 match catalog.load_table(&table_ident).await {
1110 Ok(table) => {
1111 let current_schema = table.metadata().current_schema();
1114 if let Some(mismatch) = schema_mismatch(current_schema, schema) {
1115 anyhow::bail!(
1116 "Iceberg table '{}' schema does not match expected schema: {}",
1117 table_name,
1118 mismatch
1119 );
1120 }
1121 Ok(table)
1122 }
1123 Err(err) => {
1124 if matches!(err.kind(), ErrorKind::TableNotFound { .. })
1125 || err
1126 .message()
1127 .contains("Tried to load a table that does not exist")
1128 {
1129 let table_creation = TableCreation::builder()
1133 .name(table_name.clone())
1134 .schema(schema.clone())
1135 .build();
1139
1140 catalog
1141 .create_table(&namespace_ident, table_creation)
1142 .await
1143 .with_context(|| {
1144 format!(
1145 "Failed to create Iceberg table '{}' in namespace '{}'",
1146 table_name, namespace
1147 )
1148 })
1149 } else {
1150 Err(err).context("Failed to load Iceberg table")
1152 }
1153 }
1154 }
1155}
1156
1157fn retrieve_upper_from_snapshots(
1164 snapshots: &mut [Arc<Snapshot>],
1165) -> anyhow::Result<Option<(Antichain<Timestamp>, GlobalId, u64)>> {
1166 snapshots.sort_by(|a, b| Ord::cmp(&b.sequence_number(), &a.sequence_number()));
1167
1168 for snapshot in snapshots {
1169 let props = &snapshot.summary().additional_properties;
1170 if let (Some(frontier_json), Some(sink_id_str), Some(sink_version_str)) = (
1171 props.get("mz-frontier"),
1172 props.get("mz-sink-id"),
1173 props.get("mz-sink-version"),
1174 ) {
1175 let frontier: Vec<Timestamp> = serde_json::from_str(frontier_json)
1176 .context("Failed to deserialize frontier from snapshot properties")?;
1177 let frontier = Antichain::from_iter(frontier);
1178
1179 let sink_id = sink_id_str
1180 .parse::<GlobalId>()
1181 .context("Failed to parse mz-sink-id from snapshot properties")?;
1182
1183 let sink_version = sink_version_str
1184 .parse::<u64>()
1185 .context("Failed to parse mz-sink-version from snapshot properties")?;
1186
1187 return Ok(Some((frontier, sink_id, sink_version)));
1188 }
1189 if snapshot.summary().operation.as_str() != "replace" {
1190 anyhow::bail!(
1195 "Iceberg table is in an inconsistent state: snapshot {} has operation '{}' but is missing 'mz-frontier' property. Schema or partition spec evolution is not supported.",
1196 snapshot.snapshot_id(),
1197 snapshot.summary().operation.as_str(),
1198 );
1199 }
1200 }
1201
1202 Ok(None)
1203}
1204
1205fn relation_desc_to_iceberg_schema(
1215 desc: &mz_repr::RelationDesc,
1216) -> anyhow::Result<(ArrowSchema, SchemaRef)> {
1217 let arrow_schema =
1218 mz_arrow_util::builder::desc_to_schema_with_overrides(desc, iceberg_type_overrides)
1219 .context("Failed to convert RelationDesc to Iceberg-compatible Arrow schema")?;
1220
1221 let arrow_schema_with_ids = add_field_ids_to_arrow_schema(arrow_schema);
1222
1223 let iceberg_schema = arrow_schema_to_schema_auto_assign_ids(&arrow_schema_with_ids)
1224 .context("Failed to convert Arrow schema to Iceberg schema")?;
1225
1226 Ok((arrow_schema_with_ids, Arc::new(iceberg_schema)))
1227}
1228
1229fn equality_ids_for_indices(
1234 current_schema: &Schema,
1235 materialize_arrow_schema: &ArrowSchema,
1236 equality_indices: &[usize],
1237) -> anyhow::Result<Vec<i32>> {
1238 let top_level_fields = current_schema.as_struct();
1239
1240 equality_indices
1241 .iter()
1242 .map(|index| {
1243 let mz_field = materialize_arrow_schema
1244 .fields()
1245 .get(*index)
1246 .with_context(|| format!("Equality delete key index {index} is out of bounds"))?;
1247 let field_name = mz_field.name();
1248 let iceberg_field = top_level_fields
1249 .field_by_name(field_name)
1250 .with_context(|| {
1251 format!(
1252 "Equality delete key column '{}' not found in Iceberg table schema",
1253 field_name
1254 )
1255 })?;
1256 Ok(iceberg_field.id)
1257 })
1258 .collect()
1259}
1260
1261fn build_schema_with_op_column(schema: &ArrowSchema) -> ArrowSchema {
1263 let mut fields: Vec<Arc<Field>> = schema.fields().iter().cloned().collect();
1264 fields.push(Arc::new(Field::new("__op", DataType::Int32, false)));
1265 ArrowSchema::new(fields)
1266}
1267
1268#[allow(clippy::disallowed_types)]
1273fn build_schema_with_append_columns(schema: &ArrowSchema) -> ArrowSchema {
1274 use mz_storage_types::sinks::{ICEBERG_APPEND_DIFF_COLUMN, ICEBERG_APPEND_TIMESTAMP_COLUMN};
1275 let mut fields: Vec<Arc<Field>> = schema.fields().iter().cloned().collect();
1276 fields.push(Arc::new(Field::new(
1277 ICEBERG_APPEND_DIFF_COLUMN,
1278 DataType::Int32,
1279 false,
1280 )));
1281 fields.push(Arc::new(Field::new(
1282 ICEBERG_APPEND_TIMESTAMP_COLUMN,
1283 DataType::Int64,
1284 false,
1285 )));
1286
1287 add_field_ids_to_arrow_schema(ArrowSchema::new(fields).with_metadata(schema.metadata().clone()))
1288}
1289
1290fn mint_batch_descriptions<'scope>(
1295 name: String,
1296 sink_id: GlobalId,
1297 input: SinkBatchStream<'scope>,
1298 sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
1299 connection: IcebergSinkConnection,
1300 storage_configuration: StorageConfiguration,
1301 initial_schema: SchemaRef,
1302) -> (
1303 StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
1304 StreamVec<'scope, Timestamp, Infallible>,
1305 StreamVec<'scope, Timestamp, HealthStatusMessage>,
1306 PressOnDropButton,
1307) {
1308 let scope = input.scope();
1309 let name_for_error = name.clone();
1310 let name_for_logging = name.clone();
1311 let mut builder = OperatorBuilder::new(name, scope.clone());
1312 let sink_version = sink.version;
1313
1314 let hashed_id = sink_id.hashed();
1315 let is_active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
1316 let (_, table_ready_stream) = builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
1317 let (batch_desc_output, batch_desc_stream) =
1318 builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
1319 let mut input = builder.new_input_for(input, Pipeline, &batch_desc_output);
1320
1321 let as_of = sink.as_of.clone();
1322 let commit_interval = sink
1323 .commit_interval
1324 .expect("the planner should have enforced this")
1325 .clone();
1326
1327 let (button, errors): (_, StreamVec<'scope, Timestamp, Rc<anyhow::Error>>) =
1328 builder.build_fallible(move |caps| {
1329 Box::pin(async move {
1330 let [table_ready_capset, capset]: &mut [_; 2] = caps.try_into().unwrap();
1331
1332 if !is_active_worker {
1333 return Ok(());
1335 }
1336
1337 let table_ident = TableIdent::new(
1338 NamespaceIdent::new(connection.namespace.clone()),
1339 connection.table.clone(),
1340 );
1341 let catalog = connection
1342 .catalog_connection
1343 .connect(&storage_configuration, InTask::Yes, Some(&table_ident))
1344 .await
1345 .with_context(|| {
1346 format!(
1347 "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
1348 connection.catalog_connection.uri, connection.namespace, connection.table
1349 )
1350 })?;
1351
1352 let table = load_or_create_table(
1353 catalog.as_ref(),
1354 connection.namespace.clone(),
1355 connection.table.clone(),
1356 initial_schema.as_ref(),
1357 )
1358 .await?;
1359 debug!(
1360 ?sink_id,
1361 %name_for_logging,
1362 namespace = %connection.namespace,
1363 table = %connection.table,
1364 "iceberg mint loaded/created table"
1365 );
1366
1367 *table_ready_capset = CapabilitySet::new();
1368
1369 let mut snapshots: Vec<_> = table.metadata().snapshots().cloned().collect();
1370 let resume = retrieve_upper_from_snapshots(&mut snapshots)?;
1371 let (resume_upper, resume_version) = match resume {
1372 Some((f, _, v)) => (f, v),
1373 None => (Antichain::from_elem(Timestamp::minimum()), 0),
1374 };
1375 debug!(
1376 ?sink_id,
1377 %name_for_logging,
1378 resume_upper = %resume_upper.pretty(),
1379 resume_version,
1380 as_of = %as_of.pretty(),
1381 "iceberg mint resume position loaded"
1382 );
1383
1384 let overcompacted =
1386 *resume_upper != [Timestamp::minimum()] &&
1388 PartialOrder::less_than(&resume_upper, &as_of);
1390
1391 if overcompacted {
1392 let err = super::overcompacted_error(
1395 &name_for_error,
1396 &as_of,
1397 &resume_upper,
1398 "drop the Iceberg table and recreate the sink",
1399 );
1400 return Err(anyhow::anyhow!("{err}"));
1404 };
1405
1406 if resume_version > sink_version {
1407 anyhow::bail!("Fenced off by newer sink version: resume_version {}, sink_version {}", resume_version, sink_version);
1408 }
1409
1410 let mut initialized = false;
1411 let mut observed_frontier;
1412 let mut minted_batches = VecDeque::new();
1417
1418 let catchup_start = if *resume_upper == [Timestamp::minimum()] {
1421 let batch_upper = Antichain::from_elem(
1423 as_of.as_option().expect("as_of not empty").step_forward());
1424 let batch = (as_of.clone(), batch_upper.clone());
1425 minted_batches.push_back(batch.clone());
1426 batch_desc_output.give(&capset[0], batch);
1427 capset.downgrade(batch_upper.clone());
1428
1429 batch_upper
1431 } else {
1432 resume_upper.clone()
1434 };
1435
1436 loop {
1437 if let Some(event) = input.next().await {
1438 match event {
1439 Event::Data(_, _) => continue,
1440 Event::Progress(frontier) => {
1441 observed_frontier = frontier;
1442 }
1443 }
1444 } else {
1445 return Ok(());
1446 }
1447
1448 if !initialized {
1449 if observed_frontier.is_empty() {
1450 if catchup_start.is_empty() {
1457 return Ok(());
1460 }
1461 debug!(
1462 ?sink_id,
1463 %name_for_logging,
1464 batch_lower = %catchup_start.pretty(),
1465 "iceberg mint input closed before initialization; minting final batch"
1466 );
1467 let batch = (catchup_start.clone(), Antichain::new());
1468 batch_desc_output.give(&capset[0], batch);
1469 return Ok(());
1470 }
1471
1472 if !PartialOrder::less_than(&catchup_start, &observed_frontier)
1475 {
1476 continue;
1477 }
1478
1479 let mut batch_descriptions = vec![];
1480 let mut current_upper = observed_frontier.clone();
1481 let current_upper_ts = observed_frontier.as_option().expect("frontier not empty").clone();
1482 debug!(
1483 ?sink_id,
1484 %name_for_logging,
1485 batch_lower = %catchup_start.pretty(),
1486 current_upper = %current_upper.pretty(),
1487 "iceberg mint initializing (catch-up batch)"
1488 );
1489 debug!(
1490 "{}: creating catch-up batch [{}, {})",
1491 name_for_logging,
1492 catchup_start.pretty(),
1493 current_upper.pretty()
1494 );
1495 batch_descriptions.push((catchup_start.clone(), current_upper.clone()));
1496
1497 for i in 1..INITIAL_DESCRIPTIONS_TO_MINT {
1499 let duration_millis = commit_interval.as_millis()
1500 .checked_mul(u128::from(i))
1501 .expect("commit interval multiplication overflow");
1502 let duration_ts = Timestamp::new(
1503 u64::try_from(duration_millis)
1504 .expect("commit interval too large for u64"),
1505 );
1506 let desired_batch_upper = Antichain::from_elem(
1507 current_upper_ts.step_forward_by(&duration_ts),
1508 );
1509
1510 let batch_description =
1511 (current_upper.clone(), desired_batch_upper.clone());
1512 debug!(
1513 "{}: minting future batch {}/{} [{}, {})",
1514 name_for_logging,
1515 i,
1516 INITIAL_DESCRIPTIONS_TO_MINT,
1517 current_upper.pretty(),
1518 desired_batch_upper.pretty()
1519 );
1520 current_upper = batch_description.1.clone();
1521 batch_descriptions.push(batch_description);
1522 }
1523
1524 minted_batches.extend(batch_descriptions.clone());
1525
1526 for desc in batch_descriptions {
1527 batch_desc_output.give(&capset[0], desc);
1528 }
1529
1530 capset.downgrade(current_upper);
1531
1532 initialized = true;
1533 } else {
1534 if observed_frontier.is_empty() {
1535 return Ok(());
1537 }
1538 while let Some(oldest_desc) = minted_batches.front() {
1541 let oldest_upper = &oldest_desc.1;
1542 if !PartialOrder::less_equal(oldest_upper, &observed_frontier) {
1543 break;
1544 }
1545
1546 let newest_upper = minted_batches.back().unwrap().1.clone();
1547 let new_lower = newest_upper.clone();
1548 let duration_ts = Timestamp::new(commit_interval.as_millis()
1549 .try_into()
1550 .expect("commit interval too large for u64"));
1551 let new_upper = Antichain::from_elem(newest_upper
1552 .as_option()
1553 .unwrap()
1554 .step_forward_by(&duration_ts));
1555
1556 let new_batch_description = (new_lower.clone(), new_upper.clone());
1557 minted_batches.pop_front();
1558 minted_batches.push_back(new_batch_description.clone());
1559
1560 batch_desc_output.give(&capset[0], new_batch_description);
1561
1562 capset.downgrade(new_upper);
1563 }
1564 }
1565 }
1566 })
1567 });
1568
1569 let statuses = errors.map(|error| HealthStatusMessage {
1570 id: None,
1571 update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
1572 namespace: StatusNamespace::Iceberg,
1573 });
1574 (
1575 batch_desc_stream,
1576 table_ready_stream,
1577 statuses,
1578 button.press_on_drop(),
1579 )
1580}
1581
1582#[derive(Clone, Debug, Serialize, Deserialize)]
1583#[serde(try_from = "AvroDataFile", into = "AvroDataFile")]
1584struct SerializableDataFile {
1585 pub data_file: DataFile,
1586 pub schema: Schema,
1587}
1588
1589#[derive(Clone, Debug, Serialize, Deserialize)]
1597struct AvroDataFile {
1598 pub data_file: Vec<u8>,
1599 pub schema: Vec<u8>,
1601}
1602
1603impl From<SerializableDataFile> for AvroDataFile {
1604 fn from(value: SerializableDataFile) -> Self {
1605 let mut data_file = Vec::new();
1606 write_data_files_to_avro(
1607 &mut data_file,
1608 [value.data_file],
1609 &StructType::new(vec![]),
1610 FormatVersion::V2,
1611 )
1612 .expect("serialization into buffer");
1613 let schema = serde_json::to_vec(&value.schema).expect("schema serialization");
1614 AvroDataFile { data_file, schema }
1615 }
1616}
1617
1618impl TryFrom<AvroDataFile> for SerializableDataFile {
1619 type Error = String;
1620
1621 fn try_from(value: AvroDataFile) -> Result<Self, Self::Error> {
1622 let schema: Schema = serde_json::from_slice(&value.schema)
1623 .map_err(|e| format!("Failed to deserialize schema: {}", e))?;
1624 let data_files = read_data_files_from_avro(
1625 &mut &*value.data_file,
1626 &schema,
1627 0,
1628 &StructType::new(vec![]),
1629 FormatVersion::V2,
1630 )
1631 .map_err_to_string_with_causes()?;
1632 let Some(data_file) = data_files.into_iter().next() else {
1633 return Err("No DataFile found in Avro data".into());
1634 };
1635 Ok(SerializableDataFile { data_file, schema })
1636 }
1637}
1638
1639#[derive(Clone, Debug, Serialize, Deserialize)]
1641struct BoundedDataFile {
1642 pub data_file: SerializableDataFile,
1643 pub batch_desc: (Antichain<Timestamp>, Antichain<Timestamp>),
1644}
1645
1646impl BoundedDataFile {
1647 pub fn new(
1648 file: DataFile,
1649 schema: Schema,
1650 batch_desc: (Antichain<Timestamp>, Antichain<Timestamp>),
1651 ) -> Self {
1652 Self {
1653 data_file: SerializableDataFile {
1654 data_file: file,
1655 schema,
1656 },
1657 batch_desc,
1658 }
1659 }
1660
1661 pub fn batch_desc(&self) -> &(Antichain<Timestamp>, Antichain<Timestamp>) {
1662 &self.batch_desc
1663 }
1664
1665 pub fn data_file(&self) -> &DataFile {
1666 &self.data_file.data_file
1667 }
1668
1669 pub fn into_data_file(self) -> DataFile {
1670 self.data_file.data_file
1671 }
1672}
1673
1674#[derive(Clone, Debug, Default)]
1676struct BoundedDataFileSet {
1677 pub data_files: Vec<BoundedDataFile>,
1678}
1679
1680fn data_file_location(configured_path: Option<&str>, location: &str) -> String {
1690 if let Some(path) = configured_path {
1697 return path.trim_end_matches('/').to_string();
1698 }
1699
1700 let corrected_location = match location.rsplit_once("/metadata/") {
1704 Some((a, b)) if b.ends_with(".metadata.json") => a,
1705 _ => location,
1706 };
1707 format!("{}/data", corrected_location.trim_end_matches('/'))
1710}
1711
1712fn write_data_files<'scope, H: EnvelopeHandler + 'static>(
1718 name: String,
1719 input: SinkBatchStream<'scope>,
1720 batch_desc_input: StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
1721 table_ready_stream: StreamVec<'scope, Timestamp, Infallible>,
1722 sink_id: GlobalId,
1723 from_id: GlobalId,
1724 key_is_synthetic: bool,
1725 as_of: Antichain<Timestamp>,
1726 connection: IcebergSinkConnection,
1727 storage_configuration: StorageConfiguration,
1728 materialize_arrow_schema: Arc<ArrowSchema>,
1729 metrics: Arc<IcebergSinkMetrics>,
1730 statistics: SinkStatistics,
1731) -> (
1732 StreamVec<'scope, Timestamp, BoundedDataFile>,
1733 StreamVec<'scope, Timestamp, HealthStatusMessage>,
1734 PressOnDropButton,
1735) {
1736 let scope = input.scope();
1737 let name_for_logging = name.clone();
1738 let mut builder = OperatorBuilder::new(name, scope.clone());
1739
1740 let (output, output_stream) = builder.new_output::<CapacityContainerBuilder<_>>();
1741
1742 let mut table_ready_input = builder.new_disconnected_input(table_ready_stream, Pipeline);
1743 let mut batch_desc_input =
1744 builder.new_input_for(batch_desc_input.broadcast(), Pipeline, &output);
1745 let mut input = builder.new_disconnected_input(input, Pipeline);
1746
1747 let (button, errors): (_, StreamVec<'scope, Timestamp, Rc<anyhow::Error>>) = builder
1748 .build_fallible(move |caps| {
1749 Box::pin(async move {
1750 let [capset]: &mut [_; 1] = caps.try_into().unwrap();
1751 let namespace_ident = NamespaceIdent::new(connection.namespace.clone());
1752 let table_ident = TableIdent::new(namespace_ident, connection.table.clone());
1753 let catalog = connection
1754 .catalog_connection
1755 .connect(&storage_configuration, InTask::Yes, Some(&table_ident))
1756 .await
1757 .with_context(|| {
1758 format!(
1759 "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
1760 connection.catalog_connection.uri,
1761 connection.namespace,
1762 connection.table
1763 )
1764 })?;
1765
1766 while let Some(_) = table_ready_input.next().await {
1767 }
1769 let table = catalog.load_table(&table_ident).await.with_context(|| {
1770 format!(
1771 "Failed to load Iceberg table '{}.{}' in write_data_files operator",
1772 connection.namespace, connection.table
1773 )
1774 })?;
1775
1776 let table_metadata = table.metadata().clone();
1777 let current_schema = Arc::clone(table_metadata.current_schema());
1778
1779 let arrow_schema = Arc::new(
1783 merge_materialize_metadata_into_iceberg_schema(
1784 materialize_arrow_schema.as_ref(),
1785 current_schema.as_ref(),
1786 )
1787 .context("Failed to merge Materialize metadata into Iceberg schema")?,
1788 );
1789
1790 let properties = table_metadata.properties();
1800 let configured_path = properties
1801 .get("write.data.path")
1802 .or_else(|| properties.get("write.folder-storage.path"));
1803 let data_location = data_file_location(
1804 configured_path.map(String::as_str),
1805 table_metadata.location(),
1806 );
1807 debug!(%data_location, "iceberg sink data file location");
1808 let location_generator =
1809 DefaultLocationGenerator::with_data_location(data_location);
1810
1811 let unique_suffix = format!("-{}", uuid::Uuid::new_v4());
1813 let file_name_generator = DefaultFileNameGenerator::new(
1814 PARQUET_FILE_PREFIX.to_string(),
1815 Some(unique_suffix),
1816 iceberg::spec::DataFileFormat::Parquet,
1817 );
1818
1819 let file_io = table.file_io().clone();
1820
1821 let writer_properties = WriterProperties::new();
1822
1823 let ctx = WriterContext {
1824 arrow_schema,
1825 current_schema: Arc::clone(¤t_schema),
1826 file_io,
1827 location_generator,
1828 file_name_generator,
1829 writer_properties,
1830 };
1831 let handler = H::new(ctx, &connection, &materialize_arrow_schema)?;
1832 let mut pk_warner =
1833 (!key_is_synthetic).then(|| PkViolationWarner::new(sink_id, from_id));
1834
1835 let mut stashed_rows: VecDeque<Span<Timestamp, ArcBatch<OrdValBatch<_>>>> =
1839 VecDeque::new();
1840
1841 let mut in_flight_batches: VecDeque<(
1845 (Antichain<Timestamp>, Antichain<Timestamp>),
1846 Box<dyn IcebergWriter>,
1847 )> = VecDeque::new();
1848
1849 let mut last_batch_desc: Option<BatchDescription> = None;
1853 let mut last_input_bounds: Option<(Antichain<Timestamp>, Antichain<Timestamp>)> =
1854 None;
1855
1856 let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
1857 let mut input_frontier = Antichain::from_elem(Timestamp::minimum());
1858
1859 while !(batch_description_frontier.is_empty() && input_frontier.is_empty()) {
1860 tokio::select! {
1861 _ = batch_desc_input.ready() => {},
1862 _ = input.ready() => {}
1863 }
1864
1865 while let Some(event) = batch_desc_input.next_sync() {
1869 match event {
1870 Event::Data(_cap, data) => {
1871 for batch_desc in data {
1872 let (lower, upper) = &batch_desc;
1873
1874 if let Some((prev_lower, prev_upper)) = last_batch_desc.as_ref()
1875 {
1876 if prev_upper != lower {
1877 anyhow::bail!(
1878 "batch descriptions must arrive in order, non-overlapping, \
1879 and without gaps: previous [{}, {}), new [{}, {})",
1880 prev_lower.pretty(),
1881 prev_upper.pretty(),
1882 lower.pretty(),
1883 upper.pretty(),
1884 );
1885 }
1886 }
1887 last_batch_desc = Some(batch_desc.clone());
1888
1889 let is_snapshot = lower == &as_of;
1891 debug!(
1892 "{}: received batch description [{}, {}), snapshot={}",
1893 name_for_logging,
1894 lower.pretty(),
1895 upper.pretty(),
1896 is_snapshot
1897 );
1898 let batch_writer = handler.create_writer(is_snapshot).await?;
1899 in_flight_batches.push_back((batch_desc.clone(), batch_writer));
1900 }
1901 }
1902 Event::Progress(frontier) => {
1903 batch_description_frontier = frontier;
1904 }
1905 }
1906 }
1907
1908 while let Some(event) = input.next_sync() {
1910 match event {
1911 Event::Data(_cap, data) => {
1912 for rows in &data {
1913 if let Some((prev_lower, prev_upper)) =
1914 last_input_bounds.as_ref()
1915 {
1916 if !PartialOrder::less_equal(prev_upper, rows.lower()) {
1920 anyhow::bail!(
1921 "input batches must arrive in order and \
1922 non-overlapping: previous [{}, {}), new [{}, {})",
1923 prev_lower.pretty(),
1924 prev_upper.pretty(),
1925 rows.lower().pretty(),
1926 rows.upper().pretty(),
1927 );
1928 }
1929 }
1930 last_input_bounds =
1931 Some((rows.lower().clone(), rows.upper().clone()));
1932
1933 stashed_rows.push_back(rows.clone());
1934 }
1935 }
1936 Event::Progress(frontier) => {
1937 input_frontier = frontier;
1938 }
1939 }
1940 }
1941
1942 metrics.stashed_rows.set(u64::cast_from(
1943 stashed_rows
1944 .iter()
1945 .map(|rows| rows.inner.as_ref().map_or(0, |b| b.len()))
1946 .sum::<usize>(),
1947 ));
1948
1949 let mut staged_messages_since_flush: u64 = 0;
1954
1955 let write_rows = async |rows: &Span<Timestamp, ArcBatch<OrdValBatch<_>>>,
1957 (lower, upper): BatchDescription,
1958 batch_writer: &mut Box<dyn IcebergWriter>|
1959 -> Result<(), anyhow::Error> {
1960 let Some(rows) = &rows.inner else {
1962 return Ok(());
1963 };
1964 for_each_diff_pair_async(
1965 rows,
1966 Some(lower),
1967 Some(upper),
1968 async |key, time, diff_pair| -> Result<(), anyhow::Error> {
1969 if let Some(warner) = pk_warner.as_mut() {
1970 warner.observe(key, time);
1971 }
1972
1973 let record_batch = handler
1974 .row_to_batch(diff_pair, time)
1975 .context("failed to convert row to recordbatch")?;
1976 staged_messages_since_flush +=
1977 u64::cast_from(record_batch.num_rows());
1978 batch_writer
1979 .write(record_batch)
1980 .await
1981 .context("failed to write recordbatch")?;
1982 if staged_messages_since_flush >= 10_000 {
1983 statistics.inc_messages_staged_by(staged_messages_since_flush);
1984 staged_messages_since_flush = 0;
1985 }
1986 Ok(())
1987 },
1988 )
1989 .await?;
1990 if let Some(warner) = pk_warner.as_mut() {
1994 warner.flush();
1995 }
1996 Ok(())
1997 };
1998
1999 let close_batch = async |batch_desc: BatchDescription,
2001 batch_writer: &mut Box<dyn IcebergWriter>|
2002 -> Result<(), anyhow::Error> {
2003 let close_started_at = Instant::now();
2004 let data_files = batch_writer.close().await;
2005 metrics
2006 .writer_close_duration_seconds
2007 .observe(close_started_at.elapsed().as_secs_f64());
2008 let data_files = data_files.context("Failed to close batch writer")?;
2009 debug!(
2010 "{}: closed batch [{}, {}), wrote {} files",
2011 name_for_logging,
2012 batch_desc.0.pretty(),
2013 batch_desc.1.pretty(),
2014 data_files.len()
2015 );
2016 for data_file in data_files {
2017 match data_file.content_type() {
2018 iceberg::spec::DataContentType::Data => {
2019 metrics.data_files_written.inc();
2020 }
2021 iceberg::spec::DataContentType::PositionDeletes
2022 | iceberg::spec::DataContentType::EqualityDeletes => {
2023 metrics.delete_files_written.inc();
2024 }
2025 }
2026 statistics.inc_bytes_staged_by(data_file.file_size_in_bytes());
2027 let file = BoundedDataFile::new(
2028 data_file,
2029 current_schema.as_ref().clone(),
2030 batch_desc.clone(),
2031 );
2032 output.give(&capset[0], file);
2033 }
2034
2035 capset.downgrade(batch_desc.1.clone());
2038 Ok(())
2039 };
2040
2041 with_ready_batches(
2043 input_frontier.clone(),
2044 &mut stashed_rows,
2045 batch_description_frontier.clone(),
2046 &mut in_flight_batches,
2047 write_rows,
2048 close_batch,
2049 )
2050 .await?;
2051
2052 if staged_messages_since_flush > 0 {
2053 statistics.inc_messages_staged_by(staged_messages_since_flush);
2054 }
2055 metrics.stashed_rows.set(u64::cast_from(
2056 stashed_rows
2057 .iter()
2058 .map(|rows| rows.inner.as_ref().map_or(0, |b| b.len()))
2059 .sum::<usize>(),
2060 ));
2061 }
2062 Ok(())
2063 })
2064 });
2065
2066 let statuses = errors.map(|error| HealthStatusMessage {
2067 id: None,
2068 update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
2069 namespace: StatusNamespace::Iceberg,
2070 });
2071 (output_stream, statuses, button.press_on_drop())
2072}
2073
2074type BatchDescription = (Antichain<Timestamp>, Antichain<Timestamp>);
2076
2077async fn with_ready_batches<L: Layout, W, Write, Close>(
2089 input_frontier: Antichain<Timestamp>,
2090 input_batches: &mut VecDeque<Span<Timestamp, ArcBatch<OrdValBatch<L>>>>,
2091 output_frontier: Antichain<Timestamp>,
2092 output_batches: &mut VecDeque<(BatchDescription, W)>,
2093 mut write_rows: Write,
2094 mut close_batch: Close,
2095) -> Result<(), anyhow::Error>
2096where
2097 L::TimeContainer: BatchContainer<Owned = Timestamp>,
2098 Write: AsyncFnMut(
2099 &Span<Timestamp, ArcBatch<OrdValBatch<L>>>,
2100 BatchDescription,
2101 &mut W,
2102 ) -> Result<(), anyhow::Error>,
2103 Close: AsyncFnMut(BatchDescription, &mut W) -> Result<(), anyhow::Error>,
2104{
2105 loop {
2106 {
2107 let output_lower = output_batches
2110 .front()
2111 .map_or(&output_frontier, |((lower, _), _)| lower);
2112 while input_batches
2113 .pop_front_if(|rows| PartialOrder::less_equal(rows.upper(), output_lower))
2114 .is_some()
2115 {}
2116 }
2117
2118 {
2119 let input_lower = input_batches
2122 .front()
2123 .map_or(&input_frontier, |rows| rows.lower());
2124 while let Some((batch_desc, mut batch_writer)) =
2125 output_batches.pop_front_if(|((_, batch_upper), _)| {
2126 PartialOrder::less_equal(batch_upper, input_lower)
2127 })
2128 {
2129 close_batch(batch_desc, &mut batch_writer).await?;
2130 }
2131 }
2132
2133 let Some((batch_desc, batch_writer)) = output_batches.front_mut() else {
2134 break;
2136 };
2137
2138 let Some(rows) = input_batches.front() else {
2139 break;
2141 };
2142
2143 write_rows(rows, batch_desc.clone(), batch_writer).await?;
2151 let output_upper = batch_desc.1.clone();
2152 let rows_upper = rows.upper();
2153 if PartialOrder::less_equal(&output_upper, rows_upper) {
2154 let (batch_desc, mut batch_writer) =
2156 output_batches.pop_front().expect("already checked front");
2157 close_batch(batch_desc, &mut batch_writer).await?;
2158 }
2159 if PartialOrder::less_equal(rows_upper, &output_upper) {
2160 input_batches.pop_front();
2162 }
2163
2164 }
2168
2169 Ok(())
2170}
2171
2172#[cfg(test)]
2173mod tests {
2174 use iceberg::spec::{ListType, MapType, PrimitiveType, Type};
2175 use iceberg::writer::file_writer::location_generator::LocationGenerator;
2176 use mz_repr::SqlScalarType;
2177 use mz_storage_types::sinks::ICEBERG_UINT64_DECIMAL_PRECISION;
2178
2179 use super::*;
2180
2181 fn manifest_uri(configured_path: Option<&str>, location: &str) -> String {
2183 let data_location = data_file_location(configured_path, location);
2184 DefaultLocationGenerator::with_data_location(data_location)
2185 .generate_location(None, "part-00000.parquet")
2186 }
2187
2188 fn assert_addresses_one_object(uri: &str) {
2192 let path = uri
2193 .split_once("://")
2194 .map(|(_scheme, path)| path)
2195 .unwrap_or(uri);
2196 assert!(
2197 !path.contains("//"),
2198 "URI has an empty path segment, so it does not name the object written: {uri}"
2199 );
2200 }
2201
2202 #[mz_ore::test]
2203 fn test_data_file_location_trims_configured_path() {
2204 assert_eq!(
2206 manifest_uri(Some("s3://bucket/tbl/data"), "s3://bucket/tbl"),
2207 "s3://bucket/tbl/data/part-00000.parquet"
2208 );
2209
2210 for configured in [
2212 "s3://bucket/tbl/data/",
2213 "s3://bucket/tbl/data//",
2214 "s3://bucket/tbl/data///",
2215 ] {
2216 let uri = manifest_uri(Some(configured), "s3://bucket/tbl");
2217 assert_addresses_one_object(&uri);
2218 assert_eq!(uri, "s3://bucket/tbl/data/part-00000.parquet");
2219 }
2220 }
2221
2222 #[mz_ore::test]
2223 fn test_data_file_location_trims_table_location() {
2224 assert_eq!(
2226 manifest_uri(None, "s3://bucket/tbl"),
2227 "s3://bucket/tbl/data/part-00000.parquet"
2228 );
2229
2230 let uri = manifest_uri(None, "s3://bucket/tbl/");
2233 assert_addresses_one_object(&uri);
2234 assert_eq!(uri, "s3://bucket/tbl/data/part-00000.parquet");
2235 }
2236
2237 #[mz_ore::test]
2238 fn test_data_file_location_corrects_s3_tables_metadata_path() {
2239 assert_eq!(
2242 data_file_location(None, "s3://bucket/tbl/metadata/00001-abc.metadata.json"),
2243 "s3://bucket/tbl/data"
2244 );
2245
2246 assert_eq!(
2248 data_file_location(None, "s3://bucket/metadata/tbl"),
2249 "s3://bucket/metadata/tbl/data"
2250 );
2251 }
2252
2253 #[mz_ore::test]
2254 fn test_iceberg_type_overrides() {
2255 let result = iceberg_type_overrides(&SqlScalarType::UInt16);
2257 assert_eq!(result.unwrap().0, DataType::Int32);
2258
2259 let result = iceberg_type_overrides(&SqlScalarType::UInt32);
2261 assert_eq!(result.unwrap().0, DataType::Int64);
2262
2263 let result = iceberg_type_overrides(&SqlScalarType::UInt64);
2265 assert_eq!(
2266 result.unwrap().0,
2267 DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
2268 );
2269
2270 let result = iceberg_type_overrides(&SqlScalarType::MzTimestamp);
2272 assert_eq!(
2273 result.unwrap().0,
2274 DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
2275 );
2276
2277 let result = iceberg_type_overrides(&SqlScalarType::Interval);
2279 assert_eq!(result.unwrap().0, DataType::LargeUtf8);
2280
2281 let result = iceberg_type_overrides(&SqlScalarType::Uuid);
2283 assert_eq!(result.unwrap().0, DataType::Utf8);
2284
2285 assert!(iceberg_type_overrides(&SqlScalarType::Int32).is_none());
2287 assert!(iceberg_type_overrides(&SqlScalarType::String).is_none());
2288 assert!(iceberg_type_overrides(&SqlScalarType::Bool).is_none());
2289 }
2290
2291 #[mz_ore::test]
2292 fn test_iceberg_schema_with_nested_uint64() {
2293 let desc = mz_repr::RelationDesc::builder()
2296 .with_column(
2297 "items",
2298 SqlScalarType::List {
2299 element_type: Box::new(SqlScalarType::UInt64),
2300 custom_id: None,
2301 }
2302 .nullable(true),
2303 )
2304 .finish();
2305
2306 let schema =
2307 mz_arrow_util::builder::desc_to_schema_with_overrides(&desc, iceberg_type_overrides)
2308 .expect("schema conversion should succeed");
2309
2310 if let DataType::List(field) = schema.field(0).data_type() {
2312 assert_eq!(
2313 field.data_type(),
2314 &DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0)
2315 );
2316 } else {
2317 panic!("Expected List type");
2318 }
2319 }
2320
2321 #[mz_ore::test]
2322 fn test_iceberg_interval_override() {
2323 let result = iceberg_type_overrides(&SqlScalarType::Interval);
2325 assert_eq!(result.unwrap().0, DataType::LargeUtf8);
2326
2327 let desc = mz_repr::RelationDesc::builder()
2329 .with_column("id", SqlScalarType::Int32.nullable(false))
2330 .with_column("dur", SqlScalarType::Interval.nullable(true))
2331 .finish();
2332
2333 let (arrow_schema, iceberg_schema) =
2334 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2335
2336 assert_eq!(arrow_schema.field(1).data_type(), &DataType::LargeUtf8);
2338
2339 let field = iceberg_schema
2341 .as_struct()
2342 .field_by_name("dur")
2343 .expect("field should exist");
2344 assert_eq!(*field.field_type, Type::Primitive(PrimitiveType::String));
2345 }
2346
2347 #[mz_ore::test]
2350 fn test_iceberg_uuid_override() {
2351 let result = iceberg_type_overrides(&SqlScalarType::Uuid);
2352 assert_eq!(result.unwrap().0, DataType::Utf8);
2353
2354 let desc = mz_repr::RelationDesc::builder()
2355 .with_column("id", SqlScalarType::Int32.nullable(false))
2356 .with_column("u", SqlScalarType::Uuid.nullable(true))
2357 .finish();
2358
2359 let (arrow_schema, iceberg_schema) =
2360 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2361
2362 assert_eq!(arrow_schema.field(1).data_type(), &DataType::Utf8);
2363
2364 let field = iceberg_schema
2365 .as_struct()
2366 .field_by_name("u")
2367 .expect("field should exist");
2368 assert_eq!(*field.field_type, Type::Primitive(PrimitiveType::String));
2369 assert_ne!(
2370 *field.field_type,
2371 Type::Primitive(PrimitiveType::Fixed(16)),
2372 "uuid must not fall back to the default FixedSizeBinary(16) mapping"
2373 );
2374 }
2375
2376 fn with_field_type(schema: &Schema, name: &str, ty: Type) -> Schema {
2379 let fields = schema.as_struct().fields().iter().map(|f| {
2380 let mut field = (**f).clone();
2381 if field.name == name {
2382 field.field_type = Box::new(ty.clone());
2383 }
2384 Arc::new(field)
2385 });
2386 Schema::builder()
2387 .with_fields(fields)
2388 .build()
2389 .expect("valid schema")
2390 }
2391
2392 #[mz_ore::test]
2395 fn test_schema_mismatch_accepts_legacy_uuid() {
2396 let desc = mz_repr::RelationDesc::builder()
2397 .with_column("id", SqlScalarType::Int32.nullable(false))
2398 .with_column("u", SqlScalarType::Uuid.nullable(true))
2399 .finish();
2400 let (_, expected) =
2401 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2402
2403 assert_eq!(schema_mismatch(&expected, &expected), None);
2404
2405 let legacy = with_field_type(&expected, "u", Type::Primitive(PrimitiveType::Fixed(16)));
2407 assert_eq!(schema_mismatch(&legacy, &expected), None);
2408
2409 assert!(schema_mismatch(&expected, &legacy).is_some());
2412
2413 let wrong_type = with_field_type(&expected, "u", Type::Primitive(PrimitiveType::Long));
2415 assert!(schema_mismatch(&wrong_type, &expected).is_some());
2416
2417 let truncated = Schema::builder()
2419 .with_fields(expected.as_struct().fields().iter().take(1).cloned())
2420 .build()
2421 .expect("valid schema");
2422 assert!(schema_mismatch(&truncated, &expected).is_some());
2423 }
2424
2425 #[mz_ore::test]
2428 fn test_schema_mismatch_accepts_legacy_uuid_in_list() {
2429 let desc = mz_repr::RelationDesc::builder()
2430 .with_column("id", SqlScalarType::Int32.nullable(false))
2431 .with_column(
2432 "us",
2433 SqlScalarType::List {
2434 element_type: Box::new(SqlScalarType::Uuid),
2435 custom_id: None,
2436 }
2437 .nullable(true),
2438 )
2439 .finish();
2440 let (_, expected) =
2441 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2442
2443 let Type::List(list) = &*expected
2444 .as_struct()
2445 .field_by_name("us")
2446 .expect("field should exist")
2447 .field_type
2448 else {
2449 panic!("expected a list type");
2450 };
2451 let mut legacy_element = (*list.element_field).clone();
2452 legacy_element.field_type = Box::new(Type::Primitive(PrimitiveType::Fixed(16)));
2453 let legacy = with_field_type(
2454 &expected,
2455 "us",
2456 Type::List(ListType {
2457 element_field: Arc::new(legacy_element),
2458 }),
2459 );
2460
2461 assert_eq!(schema_mismatch(&legacy, &expected), None);
2462 }
2463
2464 #[allow(clippy::disallowed_methods)]
2470 #[allow(clippy::match_wildcard_for_single_variants)]
2471 fn with_catalog_field_ids(schema: &Schema) -> Schema {
2472 fn renumber_fields(fields: &[Arc<NestedField>], next: &mut i32) -> Vec<Arc<NestedField>> {
2473 let ids: Vec<i32> = fields
2474 .iter()
2475 .map(|_| {
2476 let id = *next;
2477 *next += 1;
2478 id
2479 })
2480 .collect();
2481 fields
2482 .iter()
2483 .zip(ids)
2484 .map(|(f, id)| {
2485 let mut field = (**f).clone();
2486 field.id = id;
2487 field.field_type = Box::new(renumber_type(&field.field_type, next));
2488 Arc::new(field)
2489 })
2490 .collect()
2491 }
2492 fn renumber_type(ty: &Type, next: &mut i32) -> Type {
2493 match ty {
2494 Type::Struct(s) => Type::Struct(StructType::new(renumber_fields(s.fields(), next))),
2495 Type::List(l) => {
2496 let mut element = (*l.element_field).clone();
2497 element.id = *next;
2498 *next += 1;
2499 element.field_type = Box::new(renumber_type(&element.field_type, next));
2500 Type::List(ListType {
2501 element_field: Arc::new(element),
2502 })
2503 }
2504 Type::Map(m) => {
2505 let mut key = (*m.key_field).clone();
2506 key.id = *next;
2507 *next += 1;
2508 let mut value = (*m.value_field).clone();
2509 value.id = *next;
2510 *next += 1;
2511 key.field_type = Box::new(renumber_type(&key.field_type, next));
2512 value.field_type = Box::new(renumber_type(&value.field_type, next));
2513 Type::Map(MapType::new(Arc::new(key), Arc::new(value)))
2514 }
2515 primitive => primitive.clone(),
2516 }
2517 }
2518 let mut next = 1;
2519 Schema::builder()
2520 .with_fields(renumber_fields(schema.as_struct().fields(), &mut next))
2521 .build()
2522 .expect("valid schema")
2523 }
2524
2525 fn with_array_element_type(schema: &Schema, name: &str, ty: Type) -> Schema {
2528 let fields = schema.as_struct().fields().iter().map(|f| {
2529 if f.name != name {
2530 return Arc::clone(f);
2531 }
2532 let Type::Struct(array) = &*f.field_type else {
2533 panic!("{name} should be an array struct");
2534 };
2535 let inner = array.fields().iter().map(|inner| {
2536 let mut inner = (**inner).clone();
2537 if inner.name == "items" {
2538 let Type::List(list) = &*inner.field_type else {
2539 panic!("items should be a list");
2540 };
2541 let mut element = (*list.element_field).clone();
2542 element.field_type = Box::new(ty.clone());
2543 inner.field_type = Box::new(Type::List(ListType {
2544 element_field: Arc::new(element),
2545 }));
2546 }
2547 Arc::new(inner)
2548 });
2549 let mut field = (**f).clone();
2550 field.field_type = Box::new(Type::Struct(StructType::new(inner.collect())));
2551 Arc::new(field)
2552 });
2553 Schema::builder()
2554 .with_fields(fields)
2555 .build()
2556 .expect("valid schema")
2557 }
2558
2559 fn with_field_required(schema: &Schema, name: &str, required: bool) -> Schema {
2561 let fields = schema.as_struct().fields().iter().map(|f| {
2562 let mut field = (**f).clone();
2563 if field.name == name {
2564 field.required = required;
2565 }
2566 Arc::new(field)
2567 });
2568 Schema::builder()
2569 .with_fields(fields)
2570 .build()
2571 .expect("valid schema")
2572 }
2573
2574 #[mz_ore::test]
2584 fn test_schema_mismatch_accepts_catalog_assigned_field_ids() {
2585 let desc = mz_repr::RelationDesc::builder()
2586 .with_column("id", SqlScalarType::Int64.nullable(true))
2587 .with_column(
2588 "a",
2589 SqlScalarType::Array(Box::new(SqlScalarType::Int64)).nullable(true),
2590 )
2591 .with_column("w", SqlScalarType::Int64.nullable(true))
2592 .finish();
2593 let arrow_schema =
2594 mz_arrow_util::builder::desc_to_schema_with_overrides(&desc, iceberg_type_overrides)
2595 .expect("schema conversion should succeed");
2596
2597 let arrow_schema_with_ids = add_field_ids_to_arrow_schema(arrow_schema);
2598 let expected = arrow_schema_to_schema(&arrow_schema_with_ids)
2599 .expect("schema conversion should succeed");
2600 let catalog = with_catalog_field_ids(&expected);
2601
2602 let id_of = |schema: &Schema, name: &str| {
2604 schema
2605 .field_by_name(name)
2606 .unwrap_or_else(|| panic!("field {name} should exist"))
2607 .id
2608 };
2609 assert_eq!(id_of(&catalog, "id"), id_of(&expected, "id"));
2610 assert_ne!(id_of(&catalog, "a.items"), id_of(&expected, "a.items"));
2611 assert_ne!(id_of(&catalog, "w"), id_of(&expected, "w"));
2612
2613 let wrong_type = with_field_type(&catalog, "w", Type::Primitive(PrimitiveType::String));
2615 assert!(schema_mismatch(&wrong_type, &expected).is_some());
2616 let wrong_element =
2618 with_array_element_type(&catalog, "a", Type::Primitive(PrimitiveType::String));
2619 assert!(schema_mismatch(&wrong_element, &expected).is_some());
2620 let wrong_required = with_field_required(&catalog, "w", true);
2622 assert!(schema_mismatch(&wrong_required, &expected).is_some());
2623
2624 let flat = mz_repr::RelationDesc::builder()
2627 .with_column("id", SqlScalarType::Int64.nullable(true))
2628 .with_column("w", SqlScalarType::Int64.nullable(true))
2629 .finish();
2630 let (_, flat_expected) =
2631 relation_desc_to_iceberg_schema(&flat).expect("schema conversion should succeed");
2632 assert_eq!(
2633 schema_mismatch(&with_catalog_field_ids(&flat_expected), &flat_expected),
2634 None
2635 );
2636
2637 assert_eq!(
2640 schema_mismatch(&catalog, &expected),
2641 None,
2642 "a table that differs from the expected schema only in field ids must be accepted"
2643 );
2644 }
2645
2646 #[mz_ore::test]
2647 fn test_schema_mismatch_names_the_first_differing_field() {
2648 let desc = mz_repr::RelationDesc::builder()
2649 .with_column("id", SqlScalarType::Int64.nullable(true))
2650 .with_column(
2651 "a",
2652 SqlScalarType::Array(Box::new(SqlScalarType::Int64)).nullable(true),
2653 )
2654 .with_column("w", SqlScalarType::Int64.nullable(true))
2655 .finish();
2656 let (_, expected) =
2657 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2658
2659 let wrong_type = with_field_type(&expected, "w", Type::Primitive(PrimitiveType::String));
2660 assert_eq!(
2661 schema_mismatch(&wrong_type, &expected).as_deref(),
2662 Some("`w` has type string in the table but long in the sink")
2663 );
2664
2665 let wrong_element =
2666 with_array_element_type(&expected, "a", Type::Primitive(PrimitiveType::String));
2667 let mismatch = schema_mismatch(&wrong_element, &expected).expect("schemas differ");
2668 assert!(
2669 mismatch.starts_with("`a.")
2670 && mismatch.ends_with("has type string in the table but long in the sink"),
2671 "nested mismatch should name the path into the array: {mismatch}"
2672 );
2673
2674 let wrong_required = with_field_required(&expected, "w", true);
2675 assert_eq!(
2676 schema_mismatch(&wrong_required, &expected).as_deref(),
2677 Some("`w` is required in the table but optional in the sink")
2678 );
2679
2680 let truncated = Schema::builder()
2681 .with_fields(expected.as_struct().fields().iter().take(1).cloned())
2682 .build()
2683 .expect("valid schema");
2684 assert_eq!(
2685 schema_mismatch(&truncated, &expected).as_deref(),
2686 Some("the table is missing `a`")
2687 );
2688 assert_eq!(
2689 schema_mismatch(&expected, &truncated).as_deref(),
2690 Some("the table has an extra field `a`")
2691 );
2692
2693 let mut wide = mz_repr::RelationDesc::builder();
2696 for i in 0..500 {
2697 wide = wide.with_column(format!("column_{i}"), SqlScalarType::Int64.nullable(true));
2698 }
2699 let (_, wide_expected) = relation_desc_to_iceberg_schema(&wide.finish())
2700 .expect("schema conversion should succeed");
2701 let wide_current = with_field_type(
2702 &wide_expected,
2703 "column_499",
2704 Type::Primitive(PrimitiveType::String),
2705 );
2706 assert_eq!(
2707 schema_mismatch(&wide_current, &wide_expected).as_deref(),
2708 Some("`column_499` has type string in the table but long in the sink")
2709 );
2710 }
2711
2712 #[mz_ore::test]
2713 fn test_schema_mismatch_names_an_added_or_removed_field() {
2714 let schema = |columns: &[&str]| {
2715 let mut desc = mz_repr::RelationDesc::builder();
2716 for column in columns {
2717 desc = desc.with_column(*column, SqlScalarType::String.nullable(true));
2718 }
2719 let (_, schema) = relation_desc_to_iceberg_schema(&desc.finish())
2720 .expect("schema conversion should succeed");
2721 schema
2722 };
2723 let before = schema(&["name", "address"]);
2724 let after = schema(&["id", "name", "address"]);
2725
2726 assert_eq!(
2727 schema_mismatch(&before, &after).as_deref(),
2728 Some("the table is missing `id`")
2729 );
2730 assert_eq!(
2731 schema_mismatch(&after, &before).as_deref(),
2732 Some("the table has an extra field `id`")
2733 );
2734 assert_eq!(
2735 schema_mismatch(&schema(&["address", "name"]), &before).as_deref(),
2736 Some("the table has `address` where `name` is expected")
2737 );
2738 }
2739
2740 #[mz_ore::test]
2741 fn test_iceberg_range_schema() {
2742 let desc = mz_repr::RelationDesc::builder()
2744 .with_column("id", SqlScalarType::Int32.nullable(false))
2745 .with_column(
2746 "r",
2747 SqlScalarType::Range {
2748 element_type: Box::new(SqlScalarType::Int32),
2749 }
2750 .nullable(true),
2751 )
2752 .finish();
2753
2754 let (_arrow_schema, iceberg_schema) =
2755 relation_desc_to_iceberg_schema(&desc).expect("schema conversion should succeed");
2756
2757 let field = iceberg_schema
2759 .as_struct()
2760 .field_by_name("r")
2761 .expect("field should exist");
2762 assert!(
2763 matches!(&*field.field_type, Type::Struct(_)),
2764 "range should be struct, got: {:?}",
2765 field.field_type
2766 );
2767 }
2768
2769 #[mz_ore::test]
2770 fn equality_ids_follow_iceberg_field_ids() {
2771 let map_entries = Field::new(
2772 "entries",
2773 DataType::Struct(
2774 vec![
2775 Field::new("key", DataType::Utf8, false),
2776 Field::new("value", DataType::Utf8, true),
2777 ]
2778 .into(),
2779 ),
2780 false,
2781 );
2782 let materialize_arrow_schema = ArrowSchema::new(vec![
2783 Field::new("attrs", DataType::Map(Arc::new(map_entries), false), true),
2784 Field::new("key_col", DataType::Int32, false),
2785 ]);
2786 let materialize_arrow_schema = add_field_ids_to_arrow_schema(materialize_arrow_schema);
2787 let iceberg_schema = arrow_schema_to_schema(&materialize_arrow_schema)
2788 .expect("schema conversion should succeed");
2789
2790 let equality_ids =
2791 equality_ids_for_indices(&iceberg_schema, &materialize_arrow_schema, &[1])
2792 .expect("field lookup should succeed");
2793
2794 let expected_id = iceberg_schema
2795 .as_struct()
2796 .field_by_name("key_col")
2797 .expect("top-level field should exist")
2798 .id;
2799 assert_eq!(equality_ids, vec![expected_id]);
2800 assert_ne!(expected_id, 2);
2801 }
2802
2803 #[mz_ore::test]
2808 #[allow(clippy::disallowed_types)]
2809 fn merge_map_entries_preserves_value_extension_metadata() {
2810 use std::collections::HashMap;
2811
2812 let mz_value_metadata = HashMap::from([(
2813 ARROW_EXTENSION_NAME_KEY.to_string(),
2814 "materialize.v1.string".to_string(),
2815 )]);
2816 let mz_entries = Field::new(
2817 "entries",
2818 DataType::Struct(
2819 vec![
2820 Field::new("keys", DataType::Utf8, false),
2821 Field::new("values", DataType::Utf8, true).with_metadata(mz_value_metadata),
2822 ]
2823 .into(),
2824 ),
2825 false,
2826 );
2827 let mz_map = Field::new("m", DataType::Map(Arc::new(mz_entries), false), true)
2828 .with_metadata(HashMap::from([(
2829 ARROW_EXTENSION_NAME_KEY.to_string(),
2830 "materialize.v1.map".to_string(),
2831 )]));
2832
2833 let iceberg_entries = Field::new(
2834 "key_value",
2835 DataType::Struct(
2836 vec![
2837 Field::new("key", DataType::Utf8, false),
2838 Field::new("value", DataType::Utf8, true),
2839 ]
2840 .into(),
2841 ),
2842 false,
2843 );
2844 let iceberg_map = Field::new("m", DataType::Map(Arc::new(iceberg_entries), false), true);
2845
2846 let merged = merge_field_metadata_recursive(&iceberg_map, Some(&mz_map))
2847 .expect("merge should succeed");
2848
2849 let entries = match merged.data_type() {
2850 DataType::Map(entries, _) => entries.as_ref(),
2851 other => panic!("expected Map, got {other:?}"),
2852 };
2853 let entry_fields = match entries.data_type() {
2854 DataType::Struct(fields) => fields,
2855 other => panic!("expected Struct, got {other:?}"),
2856 };
2857 assert_eq!(entry_fields[0].name(), "key");
2859 assert_eq!(entry_fields[1].name(), "value");
2860 assert_eq!(
2863 entry_fields[1].metadata().get(ARROW_EXTENSION_NAME_KEY),
2864 Some(&"materialize.v1.string".to_string()),
2865 );
2866 }
2867
2868 mod with_ready_batches {
2869 use differential_dataflow::trace::Description;
2870 use differential_dataflow::trace::implementations::Vector;
2871
2872 use super::*;
2873
2874 type TestBatch = OrdValBatch<Vector<((u64, u64), Timestamp, Diff)>>;
2875
2876 fn frontier(t: Option<u64>) -> Antichain<Timestamp> {
2878 t.map_or_else(Antichain::new, |t| Antichain::from_elem(Timestamp::new(t)))
2879 }
2880
2881 fn span(lower: u64, upper: Option<u64>) -> BatchDescription {
2883 (frontier(Some(lower)), frontier(upper))
2884 }
2885
2886 fn input(lower: u64, upper: Option<u64>) -> Span<Timestamp, ArcBatch<TestBatch>> {
2889 let (lower, upper) = span(lower, upper);
2890 let since = Antichain::from_elem(Timestamp::MIN);
2891 Span::new(Description::new(lower, upper, since), None)
2892 }
2893
2894 #[derive(Debug, PartialEq)]
2895 enum Call {
2896 Write(BatchDescription, BatchDescription),
2898 Close(BatchDescription),
2899 }
2900
2901 async fn run(
2904 input_frontier: Antichain<Timestamp>,
2905 input_batches: &mut VecDeque<Span<Timestamp, ArcBatch<TestBatch>>>,
2906 output_frontier: Antichain<Timestamp>,
2907 output_batches: &mut VecDeque<(BatchDescription, ())>,
2908 ) -> Vec<Call> {
2909 let calls = RefCell::new(vec![]);
2910 with_ready_batches(
2911 input_frontier,
2912 input_batches,
2913 output_frontier,
2914 output_batches,
2915 async |rows: &Span<Timestamp, ArcBatch<TestBatch>>, desc, _writer: &mut ()| {
2916 let bounds = (rows.lower().clone(), rows.upper().clone());
2917 calls.borrow_mut().push(Call::Write(bounds, desc));
2918 Ok(())
2919 },
2920 async |desc, _writer: &mut ()| {
2921 calls.borrow_mut().push(Call::Close(desc));
2922 Ok(())
2923 },
2924 )
2925 .await
2926 .expect("test callbacks never fail");
2927 calls.into_inner()
2928 }
2929
2930 #[mz_ore::test(tokio::test)]
2931 async fn input_batch_spanning_multiple_output_batches() {
2932 let mut inputs = VecDeque::from([input(0, Some(30))]);
2933 let mut outputs = VecDeque::from([
2934 (span(0, Some(10)), ()),
2935 (span(10, Some(20)), ()),
2936 (span(20, Some(30)), ()),
2937 ]);
2938
2939 let calls = run(
2940 frontier(Some(30)),
2941 &mut inputs,
2942 frontier(Some(30)),
2943 &mut outputs,
2944 )
2945 .await;
2946
2947 assert_eq!(
2950 calls,
2951 vec![
2952 Call::Write(span(0, Some(30)), span(0, Some(10))),
2953 Call::Close(span(0, Some(10))),
2954 Call::Write(span(0, Some(30)), span(10, Some(20))),
2955 Call::Close(span(10, Some(20))),
2956 Call::Write(span(0, Some(30)), span(20, Some(30))),
2957 Call::Close(span(20, Some(30))),
2958 ]
2959 );
2960 assert!(inputs.is_empty());
2961 assert!(outputs.is_empty());
2962 }
2963
2964 #[mz_ore::test(tokio::test)]
2965 async fn output_batch_spanning_multiple_input_batches() {
2966 let mut inputs =
2967 VecDeque::from([input(0, Some(10)), input(10, Some(20)), input(20, Some(30))]);
2968 let mut outputs = VecDeque::from([(span(0, Some(30)), ())]);
2969
2970 let calls = run(
2971 frontier(Some(30)),
2972 &mut inputs,
2973 frontier(Some(30)),
2974 &mut outputs,
2975 )
2976 .await;
2977
2978 assert_eq!(
2979 calls,
2980 vec![
2981 Call::Write(span(0, Some(10)), span(0, Some(30))),
2982 Call::Write(span(10, Some(20)), span(0, Some(30))),
2983 Call::Write(span(20, Some(30)), span(0, Some(30))),
2984 Call::Close(span(0, Some(30))),
2985 ]
2986 );
2987 assert!(inputs.is_empty());
2988 assert!(outputs.is_empty());
2989 }
2990
2991 #[mz_ore::test(tokio::test)]
2992 async fn input_batch_retained_for_future_output_batches() {
2993 let mut inputs = VecDeque::from([input(0, Some(30))]);
2994 let mut outputs = VecDeque::from([(span(0, Some(10)), ())]);
2995
2996 let calls = run(
2997 frontier(Some(30)),
2998 &mut inputs,
2999 frontier(Some(10)),
3000 &mut outputs,
3001 )
3002 .await;
3003
3004 assert_eq!(
3007 calls,
3008 vec![
3009 Call::Write(span(0, Some(30)), span(0, Some(10))),
3010 Call::Close(span(0, Some(10))),
3011 ]
3012 );
3013 assert_eq!(inputs.len(), 1);
3014 assert!(outputs.is_empty());
3015 }
3016
3017 #[mz_ore::test(tokio::test)]
3018 async fn already_committed_input_batches_dropped_unwritten() {
3019 let mut inputs = VecDeque::from([input(0, Some(10)), input(10, Some(20))]);
3020 let mut outputs = VecDeque::from([(span(20, Some(30)), ())]);
3021
3022 let calls = run(
3023 frontier(Some(20)),
3024 &mut inputs,
3025 frontier(Some(30)),
3026 &mut outputs,
3027 )
3028 .await;
3029
3030 assert_eq!(calls, vec![]);
3034 assert!(inputs.is_empty());
3035 assert_eq!(outputs.len(), 1);
3036 }
3037
3038 #[mz_ore::test(tokio::test)]
3039 async fn output_batch_closes_empty_once_input_frontier_passes() {
3040 let mut outputs = VecDeque::from([(span(0, Some(10)), ())]);
3041
3042 let calls = run(
3045 frontier(Some(5)),
3046 &mut VecDeque::new(),
3047 frontier(Some(10)),
3048 &mut outputs,
3049 )
3050 .await;
3051 assert_eq!(calls, vec![]);
3052 assert_eq!(outputs.len(), 1);
3053
3054 let calls = run(
3057 frontier(Some(10)),
3058 &mut VecDeque::new(),
3059 frontier(Some(10)),
3060 &mut outputs,
3061 )
3062 .await;
3063 assert_eq!(calls, vec![Call::Close(span(0, Some(10)))]);
3064 assert!(outputs.is_empty());
3065 }
3066
3067 #[mz_ore::test(tokio::test)]
3068 async fn final_output_batch_with_empty_upper() {
3069 let mut inputs = VecDeque::from([input(20, Some(30))]);
3070 let mut outputs = VecDeque::from([(span(20, None), ())]);
3071
3072 let calls = run(
3076 frontier(Some(30)),
3077 &mut inputs,
3078 frontier(None),
3079 &mut outputs,
3080 )
3081 .await;
3082 assert_eq!(calls, vec![Call::Write(span(20, Some(30)), span(20, None))]);
3083 assert!(inputs.is_empty());
3084 assert_eq!(outputs.len(), 1);
3085
3086 let calls = run(frontier(None), &mut inputs, frontier(None), &mut outputs).await;
3087 assert_eq!(calls, vec![Call::Close(span(20, None))]);
3088 assert!(outputs.is_empty());
3089 }
3090 }
3091}
3092
3093fn commit_to_iceberg<'scope>(
3097 name: String,
3098 sink_id: GlobalId,
3099 sink_version: u64,
3100 batch_input: StreamVec<'scope, Timestamp, BoundedDataFile>,
3101 batch_desc_input: StreamVec<'scope, Timestamp, (Antichain<Timestamp>, Antichain<Timestamp>)>,
3102 table_ready_stream: StreamVec<'scope, Timestamp, Infallible>,
3103 write_frontier: Rc<RefCell<Antichain<Timestamp>>>,
3104 connection: IcebergSinkConnection,
3105 storage_configuration: StorageConfiguration,
3106 write_handle: impl Future<
3107 Output = anyhow::Result<WriteHandle<SourceData, (), Timestamp, StorageDiff>>,
3108 > + 'static,
3109 metrics: Arc<IcebergSinkMetrics>,
3110 statistics: SinkStatistics,
3111) -> (
3112 StreamVec<'scope, Timestamp, HealthStatusMessage>,
3113 PressOnDropButton,
3114) {
3115 let scope = batch_input.scope();
3116 let mut builder = OperatorBuilder::new(name, scope.clone());
3117
3118 let hashed_id = sink_id.hashed();
3119 let is_active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
3120 let name_for_logging = format!("{sink_id}-commit-to-iceberg");
3121
3122 let mut input = builder.new_disconnected_input(batch_input, Exchange::new(move |_| hashed_id));
3123 let mut batch_desc_input =
3124 builder.new_disconnected_input(batch_desc_input, Exchange::new(move |_| hashed_id));
3125 let mut table_ready_input = builder.new_disconnected_input(table_ready_stream, Pipeline);
3126
3127 let (button, errors) = builder.build_fallible(move |_caps| {
3128 Box::pin(async move {
3129 if !is_active_worker {
3130 write_frontier.borrow_mut().clear();
3131 return Ok(());
3132 }
3133
3134 let namespace_ident = NamespaceIdent::new(connection.namespace.clone());
3135 let table_ident = TableIdent::new(namespace_ident, connection.table.clone());
3136 let catalog = connection
3137 .catalog_connection
3138 .connect(&storage_configuration, InTask::Yes, Some(&table_ident))
3139 .await
3140 .with_context(|| {
3141 format!(
3142 "Failed to connect to Iceberg catalog '{}' for table '{}.{}'",
3143 connection.catalog_connection.uri, connection.namespace, connection.table
3144 )
3145 })?;
3146
3147 let mut write_handle = write_handle.await?;
3148
3149 while let Some(_) = table_ready_input.next().await {
3150 }
3152 let mut table = catalog.load_table(&table_ident).await.with_context(|| {
3153 format!(
3154 "Failed to load Iceberg table '{}.{}' in commit_to_iceberg operator",
3155 connection.namespace, connection.table
3156 )
3157 })?;
3158
3159 #[allow(clippy::disallowed_types)]
3160 let mut batch_descriptions: std::collections::HashMap<
3161 (Antichain<Timestamp>, Antichain<Timestamp>),
3162 BoundedDataFileSet,
3163 > = std::collections::HashMap::new();
3164
3165 let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
3166 let mut input_frontier = Antichain::from_elem(Timestamp::minimum());
3167
3168 while !(batch_description_frontier.is_empty() && input_frontier.is_empty()) {
3169 tokio::select! {
3170 _ = batch_desc_input.ready() => {},
3171 _ = input.ready() => {}
3172 }
3173
3174 while let Some(event) = batch_desc_input.next_sync() {
3175 match event {
3176 Event::Data(_cap, data) => {
3177 for batch_desc in data {
3178 let prev = batch_descriptions
3179 .insert(batch_desc, BoundedDataFileSet { data_files: vec![] });
3180 if let Some(prev) = prev {
3181 anyhow::bail!(
3182 "Duplicate batch description received \
3183 in commit operator: {:?}",
3184 prev
3185 );
3186 }
3187 }
3188 }
3189 Event::Progress(frontier) => {
3190 batch_description_frontier = frontier;
3191 }
3192 }
3193 }
3194
3195 let ready_events = std::iter::from_fn(|| input.next_sync()).collect_vec();
3196 for event in ready_events {
3197 match event {
3198 Event::Data(_cap, data) => {
3199 for bounded_data_file in data {
3200 let entry = batch_descriptions
3201 .entry(bounded_data_file.batch_desc().clone())
3202 .or_default();
3203 entry.data_files.push(bounded_data_file);
3204 }
3205 }
3206 Event::Progress(frontier) => {
3207 input_frontier = frontier;
3208 }
3209 }
3210 }
3211
3212 let mut done_batches: Vec<_> = batch_descriptions
3218 .keys()
3219 .filter(|(lower, _upper)| PartialOrder::less_than(lower, &input_frontier))
3220 .cloned()
3221 .collect();
3222
3223 done_batches.sort_by(|a, b| {
3225 if PartialOrder::less_than(a, b) {
3226 Ordering::Less
3227 } else if PartialOrder::less_than(b, a) {
3228 Ordering::Greater
3229 } else {
3230 Ordering::Equal
3231 }
3232 });
3233
3234 for batch in done_batches {
3235 let file_set = batch_descriptions.remove(&batch).unwrap();
3236
3237 let mut data_files = vec![];
3238 let mut delete_files = vec![];
3239 let mut total_messages: u64 = 0;
3241 let mut total_bytes: u64 = 0;
3242 for file in file_set.data_files {
3243 total_messages += file.data_file().record_count();
3244 total_bytes += file.data_file().file_size_in_bytes();
3245 match file.data_file().content_type() {
3246 iceberg::spec::DataContentType::Data => {
3247 data_files.push(file.into_data_file());
3248 }
3249 iceberg::spec::DataContentType::PositionDeletes
3250 | iceberg::spec::DataContentType::EqualityDeletes => {
3251 delete_files.push(file.into_data_file());
3252 }
3253 }
3254 }
3255
3256 debug!(
3257 ?sink_id,
3258 %name_for_logging,
3259 lower = %batch.0.pretty(),
3260 upper = %batch.1.pretty(),
3261 data_files = data_files.len(),
3262 delete_files = delete_files.len(),
3263 total_messages,
3264 total_bytes,
3265 "iceberg commit applying batch"
3266 );
3267
3268 let instant = Instant::now();
3269
3270 let frontier = batch.1.clone();
3271 let frontier_json = serde_json::to_string(&frontier.elements())
3272 .context("Failed to serialize frontier to JSON")?;
3273 let snapshot_properties = vec![
3274 ("mz-sink-id".to_string(), sink_id.to_string()),
3275 ("mz-frontier".to_string(), frontier_json),
3276 ("mz-sink-version".to_string(), sink_version.to_string()),
3277 ];
3278
3279 let (table_state, commit_result) = Retry::default()
3280 .max_tries(5)
3281 .retry_async_with_state(table, |_, table| {
3282 let snapshot_properties = snapshot_properties.clone();
3283 let data_files = data_files.clone();
3284 let delete_files = delete_files.clone();
3285 let metrics = Arc::clone(&metrics);
3286 let catalog = Arc::clone(&catalog);
3287 let conn_namespace = connection.namespace.clone();
3288 let conn_table = connection.table.clone();
3289 let batch_lower = batch.0.clone();
3290 let batch_upper = batch.1.clone();
3291 async move {
3292 try_commit_batch(
3293 table,
3294 snapshot_properties,
3295 data_files,
3296 delete_files,
3297 catalog.as_ref(),
3298 &conn_namespace,
3299 &conn_table,
3300 sink_id,
3301 sink_version,
3302 &batch_lower,
3303 &batch_upper,
3304 &metrics,
3305 )
3306 .await
3307 }
3308 })
3309 .await;
3310 let commit_result = commit_result.with_context(|| {
3311 format!(
3312 "failed to commit batch to Iceberg table '{}.{}'",
3313 connection.namespace, connection.table
3314 )
3315 });
3316 table = table_state;
3317 let duration = instant.elapsed();
3318 metrics
3319 .commit_duration_seconds
3320 .observe(duration.as_secs_f64());
3321 commit_result?;
3322
3323 debug!(
3324 ?sink_id,
3325 %name_for_logging,
3326 lower = %batch.0.pretty(),
3327 upper = %batch.1.pretty(),
3328 total_messages,
3329 total_bytes,
3330 ?duration,
3331 "iceberg commit applied batch"
3332 );
3333
3334 metrics.snapshots_committed.inc();
3335 statistics.inc_messages_committed_by(total_messages);
3336 statistics.inc_bytes_committed_by(total_bytes);
3337
3338 let mut expect_upper = write_handle.shared_upper();
3339 loop {
3340 if PartialOrder::less_equal(&frontier, &expect_upper) {
3341 break;
3343 }
3344
3345 const EMPTY: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
3346 match write_handle
3347 .compare_and_append(EMPTY, expect_upper, frontier.clone())
3348 .await
3349 .expect("valid usage")
3350 {
3351 Ok(()) => break,
3352 Err(mismatch) => {
3353 expect_upper = mismatch.current;
3354 }
3355 }
3356 }
3357 write_frontier.borrow_mut().clone_from(&frontier);
3358 }
3359 }
3360
3361 Ok(())
3362 })
3363 });
3364
3365 let statuses = errors.map(|error| HealthStatusMessage {
3366 id: None,
3367 update: HealthStatusUpdate::halting(format!("{}", error.display_with_causes()), None),
3368 namespace: StatusNamespace::Iceberg,
3369 });
3370
3371 (statuses, button.press_on_drop())
3372}
3373
3374impl<'scope> SinkRender<'scope> for IcebergSinkConnection {
3375 fn get_key_indices(&self) -> Option<&[usize]> {
3376 self.key_desc_and_indices
3377 .as_ref()
3378 .map(|(_, indices)| indices.as_slice())
3379 }
3380
3381 fn get_relation_key_indices(&self) -> Option<&[usize]> {
3382 self.relation_key_indices.as_deref()
3383 }
3384
3385 fn render_sink(
3386 &self,
3387 storage_state: &mut StorageState,
3388 sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
3389 sink_id: GlobalId,
3390 batches: SinkBatchStream<'scope>,
3391 key_is_synthetic: bool,
3392 _err_collection: VecCollection<'scope, Timestamp, DataflowError, Diff>,
3393 ) -> (
3394 StreamVec<'scope, Timestamp, HealthStatusMessage>,
3395 Vec<PressOnDropButton>,
3396 ) {
3397 let scope = batches.scope();
3398
3399 let write_handle = {
3400 let persist = Arc::clone(&storage_state.persist_clients);
3401 let shard_meta = sink.to_storage_metadata.clone();
3402 async move {
3403 let client = persist.open(shard_meta.persist_location).await?;
3404 let handle = client
3405 .open_writer(
3406 shard_meta.data_shard,
3407 Arc::new(shard_meta.relation_desc),
3408 Arc::new(UnitSchema),
3409 Diagnostics::from_purpose("sink handle"),
3410 )
3411 .await?;
3412 Ok(handle)
3413 }
3414 };
3415
3416 let write_frontier = Rc::new(RefCell::new(Antichain::from_elem(Timestamp::minimum())));
3417 storage_state
3418 .sink_write_frontiers
3419 .insert(sink_id, Rc::clone(&write_frontier));
3420
3421 let (arrow_schema_with_ids, iceberg_schema) =
3422 match (|| -> Result<(ArrowSchema, Arc<Schema>), anyhow::Error> {
3423 let (arrow_schema_with_ids, iceberg_schema) =
3424 relation_desc_to_iceberg_schema(&sink.from_desc)?;
3425
3426 Ok(if sink.envelope == SinkEnvelope::Append {
3427 let extended_arrow = build_schema_with_append_columns(&arrow_schema_with_ids);
3432 let extended_iceberg = Arc::new(
3433 arrow_schema_to_schema_auto_assign_ids(&extended_arrow)
3434 .context("Failed to build Iceberg schema with append columns")?,
3435 );
3436
3437 (extended_arrow, extended_iceberg)
3438 } else {
3439 (arrow_schema_with_ids, iceberg_schema)
3440 })
3441 })() {
3442 Ok(schemas) => schemas,
3443 Err(err) => {
3444 let error_stream = std::iter::once(HealthStatusMessage {
3445 id: None,
3446 update: HealthStatusUpdate::halting(
3447 format!("{}", err.display_with_causes()),
3448 None,
3449 ),
3450 namespace: StatusNamespace::Iceberg,
3451 })
3452 .to_stream(scope);
3453 return (error_stream, vec![]);
3454 }
3455 };
3456
3457 let metrics = Arc::new(
3458 storage_state
3459 .metrics
3460 .get_iceberg_sink_metrics(sink_id, scope.index()),
3461 );
3462
3463 let statistics = storage_state
3464 .aggregated_statistics
3465 .get_sink(&sink_id)
3466 .expect("statistics initialized")
3467 .clone();
3468
3469 let connection_for_minter = self.clone();
3470 let (batch_descriptions, table_ready, mint_status, mint_button) = mint_batch_descriptions(
3471 format!("{sink_id}-iceberg-mint"),
3472 sink_id,
3473 batches.clone(),
3474 sink,
3475 connection_for_minter,
3476 storage_state.storage_configuration.clone(),
3477 Arc::clone(&iceberg_schema),
3478 );
3479
3480 let connection_for_writer = self.clone();
3481 let (datafiles, write_status, write_button) = match sink.envelope {
3482 SinkEnvelope::Upsert => write_data_files::<UpsertEnvelopeHandler>(
3483 format!("{sink_id}-write-data-files"),
3484 batches,
3485 batch_descriptions.clone(),
3486 table_ready.clone(),
3487 sink_id,
3488 sink.from,
3489 key_is_synthetic,
3490 sink.as_of.clone(),
3491 connection_for_writer,
3492 storage_state.storage_configuration.clone(),
3493 Arc::new(arrow_schema_with_ids.clone()),
3494 Arc::clone(&metrics),
3495 statistics.clone(),
3496 ),
3497 SinkEnvelope::Append => write_data_files::<AppendEnvelopeHandler>(
3498 format!("{sink_id}-write-data-files"),
3499 batches,
3500 batch_descriptions.clone(),
3501 table_ready.clone(),
3502 sink_id,
3503 sink.from,
3504 key_is_synthetic,
3505 sink.as_of.clone(),
3506 connection_for_writer,
3507 storage_state.storage_configuration.clone(),
3508 Arc::new(arrow_schema_with_ids.clone()),
3509 Arc::clone(&metrics),
3510 statistics.clone(),
3511 ),
3512 SinkEnvelope::Debezium => {
3513 unreachable!("Iceberg sink only supports Upsert and Append envelopes")
3514 }
3515 };
3516
3517 let connection_for_committer = self.clone();
3518 let (commit_status, commit_button) = commit_to_iceberg(
3519 format!("{sink_id}-commit-to-iceberg"),
3520 sink_id,
3521 sink.version,
3522 datafiles,
3523 batch_descriptions,
3524 table_ready,
3525 Rc::clone(&write_frontier),
3526 connection_for_committer,
3527 storage_state.storage_configuration.clone(),
3528 write_handle,
3529 Arc::clone(&metrics),
3530 statistics,
3531 );
3532
3533 let running_status = Some(HealthStatusMessage {
3534 id: None,
3535 update: HealthStatusUpdate::running(),
3536 namespace: StatusNamespace::Iceberg,
3537 })
3538 .to_stream(scope);
3539
3540 let statuses =
3541 scope.concatenate([running_status, mint_status, write_status, commit_status]);
3542
3543 (statuses, vec![mint_button, write_button, commit_button])
3544 }
3545}