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