Spark 4.1: Display write metrics on SQL UI

core/src/main/java/org/apache/iceberg/BaseTable.java

@@ -27,6 +27,7 @@
 import org.apache.iceberg.io.LocationProvider;
 import org.apache.iceberg.metrics.LoggingMetricsReporter;
 import org.apache.iceberg.metrics.MetricsReporter;
+import org.apache.iceberg.metrics.MetricsReporters;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 
 /**
@@ -41,7 +42,7 @@
 public class BaseTable implements Table, HasTableOperations, Serializable {
   private final TableOperations ops;
   private final String name;
-  private final MetricsReporter reporter;
+  private MetricsReporter reporter;
 
   public BaseTable(TableOperations ops, String name) {
     this(ops, name, LoggingMetricsReporter.instance());
@@ -58,6 +59,10 @@ public MetricsReporter reporter() {
     return reporter;
   }
 
+  public void combineMetricsReporter(MetricsReporter metricsReporter) {
+    this.reporter = MetricsReporters.combine(this.reporter, metricsReporter);
+  }
+
   @Override
   public TableOperations operations() {
     return ops;

core/src/main/java/org/apache/iceberg/metrics/InMemoryMetricsReporter.java

@@ -18,6 +18,7 @@
  */
 package org.apache.iceberg.metrics;
 
+import javax.annotation.Nullable;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 
 public class InMemoryMetricsReporter implements MetricsReporter {
@@ -35,4 +36,13 @@ public ScanReport scanReport() {
         "Metrics report is not a scan report");
     return (ScanReport) metricsReport;
   }
+
+  @Nullable
+  public CommitReport commitReport() {
+    if (metricsReport != null && metricsReport instanceof CommitReport) {
+      return (CommitReport) metricsReport;
+    } else {
+      return null;
+    }
+  }
 }

spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/SparkWriteUtil.java

@@ -23,10 +23,38 @@
 import static org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE;
 
 import java.util.Arrays;
+import java.util.List;
 import org.apache.iceberg.DistributionMode;
 import org.apache.iceberg.MetadataColumns;
 import org.apache.iceberg.Table;
+import org.apache.iceberg.metrics.CommitMetricsResult;
+import org.apache.iceberg.metrics.CommitReport;
+import org.apache.iceberg.metrics.CounterResult;
+import org.apache.iceberg.metrics.InMemoryMetricsReporter;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.relocated.com.google.common.collect.ObjectArrays;
+import org.apache.iceberg.spark.source.metrics.AddedDataFiles;
+import org.apache.iceberg.spark.source.metrics.AddedDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.AddedEqualityDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.AddedEqualityDeletes;
+import org.apache.iceberg.spark.source.metrics.AddedFileSizeInBytes;
+import org.apache.iceberg.spark.source.metrics.AddedPositionalDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.AddedPositionalDeletes;
+import org.apache.iceberg.spark.source.metrics.AddedRecords;
+import org.apache.iceberg.spark.source.metrics.RemovedDataFiles;
+import org.apache.iceberg.spark.source.metrics.RemovedDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.RemovedEqualityDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.RemovedEqualityDeletes;
+import org.apache.iceberg.spark.source.metrics.RemovedFileSizeInBytes;
+import org.apache.iceberg.spark.source.metrics.RemovedPositionalDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.RemovedPositionalDeletes;
+import org.apache.iceberg.spark.source.metrics.RemovedRecords;
+import org.apache.iceberg.spark.source.metrics.TotalDataFiles;
+import org.apache.iceberg.spark.source.metrics.TotalDeleteFiles;
+import org.apache.iceberg.spark.source.metrics.TotalEqualityDeletes;
+import org.apache.iceberg.spark.source.metrics.TotalFileSizeInBytes;
+import org.apache.iceberg.spark.source.metrics.TotalPositionalDeletes;
+import org.apache.iceberg.spark.source.metrics.TotalRecords;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.SortOrderUtil;
 import org.apache.spark.sql.connector.distributions.Distribution;
@@ -36,6 +64,8 @@
 import org.apache.spark.sql.connector.expressions.NamedReference;
 import org.apache.spark.sql.connector.expressions.SortDirection;
 import org.apache.spark.sql.connector.expressions.SortOrder;
+import org.apache.spark.sql.connector.metric.CustomMetric;
+import org.apache.spark.sql.connector.metric.CustomTaskMetric;
 import org.apache.spark.sql.connector.write.RowLevelOperation.Command;
 
 /**
@@ -256,4 +286,83 @@ private static SortOrder[] orderBy(Expression... exprs) {
   private static SortOrder sort(Expression expr) {
     return Expressions.sort(expr, SortDirection.ASCENDING);
   }
+
+  public static CustomMetric[] supportedCustomMetrics() {
+    return new CustomMetric[] {
+      new AddedDataFiles(),
+      new AddedDeleteFiles(),
+      new AddedEqualityDeletes(),
+      new AddedEqualityDeleteFiles(),
+      new AddedFileSizeInBytes(),
+      new AddedPositionalDeletes(),
+      new AddedPositionalDeleteFiles(),
+      new AddedRecords(),
+      new RemovedDataFiles(),
+      new RemovedDeleteFiles(),
+      new RemovedRecords(),
+      new RemovedEqualityDeleteFiles(),
+      new RemovedEqualityDeletes(),
+      new RemovedFileSizeInBytes(),
+      new RemovedPositionalDeleteFiles(),
+      new RemovedPositionalDeletes(),
+      new TotalDataFiles(),
+      new TotalDeleteFiles(),
+      new TotalEqualityDeletes(),
+      new TotalFileSizeInBytes(),
+      new TotalPositionalDeletes(),
+      new TotalRecords()
+    };
+  }
+
+  public static CustomTaskMetric[] customTaskMetrics(InMemoryMetricsReporter metricsReporter) {
+    List<CustomTaskMetric> metrics = Lists.newArrayList();
+    if (metricsReporter != null) {
+      CommitReport commitReport = metricsReporter.commitReport();
+      if (commitReport != null) {
+        CommitMetricsResult result = commitReport.commitMetrics();
+        addValue(new AddedDataFiles(), result.addedDataFiles(), metrics);
+        addValue(new AddedDeleteFiles(), result.addedDeleteFiles(), metrics);
+        addValue(new AddedEqualityDeletes(), result.addedEqualityDeletes(), metrics);
+        addValue(new AddedEqualityDeleteFiles(), result.addedEqualityDeleteFiles(), metrics);
+        addValue(new AddedFileSizeInBytes(), result.addedFilesSizeInBytes(), metrics);
+        addValue(new AddedPositionalDeletes(), result.addedPositionalDeletes(), metrics);
+        addValue(new AddedPositionalDeleteFiles(), result.addedPositionalDeleteFiles(), metrics);
+        addValue(new AddedRecords(), result.addedRecords(), metrics);
+        addValue(new RemovedDataFiles(), result.removedDataFiles(), metrics);
+        addValue(new RemovedDeleteFiles(), result.removedDeleteFiles(), metrics);
+        addValue(new RemovedRecords(), result.removedRecords(), metrics);
+        addValue(new RemovedEqualityDeleteFiles(), result.removedEqualityDeleteFiles(), metrics);
+        addValue(new RemovedEqualityDeletes(), result.removedEqualityDeletes(), metrics);
+        addValue(new RemovedFileSizeInBytes(), result.removedFilesSizeInBytes(), metrics);
+        addValue(
+            new RemovedPositionalDeleteFiles(), result.removedPositionalDeleteFiles(), metrics);
+        addValue(new RemovedPositionalDeletes(), result.removedPositionalDeletes(), metrics);
+        addValue(new TotalDataFiles(), result.totalDataFiles(), metrics);
+        addValue(new TotalDeleteFiles(), result.totalDeleteFiles(), metrics);
+        addValue(new TotalEqualityDeletes(), result.totalEqualityDeletes(), metrics);
+        addValue(new TotalFileSizeInBytes(), result.totalFilesSizeInBytes(), metrics);
+        addValue(new TotalPositionalDeletes(), result.totalPositionalDeletes(), metrics);
+        addValue(new TotalRecords(), result.totalRecords(), metrics);
+      }
+    }
+    return metrics.toArray(new CustomTaskMetric[0]);
+  }
+
+  private static void addValue(
+      CustomMetric metric, CounterResult result, List<CustomTaskMetric> taskMetrics) {
+    if (result != null) {
+      taskMetrics.add(
+          new CustomTaskMetric() {
+            @Override
+            public String name() {
+              return metric.name();
+            }
+
+            @Override
+            public long value() {
+              return result.value();
+            }
+          });
+    }
+  }
 }

spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeletesRewrite.java

@@ -22,6 +22,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import org.apache.iceberg.BaseTable;
 import org.apache.iceberg.DeleteFile;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.MetadataColumns;
@@ -37,16 +38,20 @@
 import org.apache.iceberg.io.FileIO;
 import org.apache.iceberg.io.OutputFileFactory;
 import org.apache.iceberg.io.PartitioningDVWriter;
+import org.apache.iceberg.metrics.InMemoryMetricsReporter;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.spark.PositionDeletesRewriteCoordinator;
 import org.apache.iceberg.spark.ScanTaskSetManager;
 import org.apache.iceberg.spark.SparkWriteConf;
+import org.apache.iceberg.spark.SparkWriteUtil;
 import org.apache.iceberg.util.DeleteFileSet;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.broadcast.Broadcast;
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.metric.CustomMetric;
+import org.apache.spark.sql.connector.metric.CustomTaskMetric;
 import org.apache.spark.sql.connector.write.BatchWrite;
 import org.apache.spark.sql.connector.write.DataWriter;
 import org.apache.spark.sql.connector.write.DataWriterFactory;
@@ -80,6 +85,7 @@ public class SparkPositionDeletesRewrite implements Write {
   private final int specId;
   private final StructLike partition;
   private final Map<String, String> writeProperties;
+  private InMemoryMetricsReporter metricsReporter;
 
   /**
    * Constructs a {@link SparkPositionDeletesRewrite}.
@@ -114,13 +120,28 @@ public class SparkPositionDeletesRewrite implements Write {
     this.specId = specId;
     this.partition = partition;
     this.writeProperties = writeConf.writeProperties();
+
+    if (this.table instanceof BaseTable) {
+      this.metricsReporter = new InMemoryMetricsReporter();
+      ((BaseTable) this.table).combineMetricsReporter(metricsReporter);
+    }
   }
 
   @Override
   public BatchWrite toBatch() {
     return new PositionDeleteBatchWrite();
   }
 
+  @Override
+  public CustomTaskMetric[] reportDriverMetrics() {
+    return SparkWriteUtil.customTaskMetrics(metricsReporter);
+  }
+
+  @Override
+  public CustomMetric[] supportedCustomMetrics() {
+    return SparkWriteUtil.supportedCustomMetrics();
+  }
+
   /** {@link BatchWrite} class for rewriting position deletes files from Spark */
   class PositionDeleteBatchWrite implements BatchWrite {
 

spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWrite.java

@@ -30,6 +30,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.function.Function;
+import org.apache.iceberg.BaseTable;
 import org.apache.iceberg.ContentFile;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.DeleteFile;
@@ -66,12 +67,14 @@
 import org.apache.iceberg.io.PartitioningWriter;
 import org.apache.iceberg.io.PositionDeltaWriter;
 import org.apache.iceberg.io.WriteResult;
+import org.apache.iceberg.metrics.InMemoryMetricsReporter;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 import org.apache.iceberg.spark.CommitMetadata;
 import org.apache.iceberg.spark.SparkSchemaUtil;
 import org.apache.iceberg.spark.SparkWriteConf;
 import org.apache.iceberg.spark.SparkWriteRequirements;
+import org.apache.iceberg.spark.SparkWriteUtil;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.CharSequenceSet;
 import org.apache.iceberg.util.DeleteFileSet;
@@ -83,6 +86,8 @@
 import org.apache.spark.sql.catalyst.expressions.JoinedRow;
 import org.apache.spark.sql.connector.distributions.Distribution;
 import org.apache.spark.sql.connector.expressions.SortOrder;
+import org.apache.spark.sql.connector.metric.CustomMetric;
+import org.apache.spark.sql.connector.metric.CustomTaskMetric;
 import org.apache.spark.sql.connector.write.DeltaBatchWrite;
 import org.apache.spark.sql.connector.write.DeltaWrite;
 import org.apache.spark.sql.connector.write.DeltaWriter;
@@ -116,6 +121,7 @@ class SparkPositionDeltaWrite implements DeltaWrite, RequiresDistributionAndOrde
   private final Map<String, String> writeProperties;
 
   private boolean cleanupOnAbort = false;
+  private InMemoryMetricsReporter metricsReporter;
 
   SparkPositionDeltaWrite(
       SparkSession spark,
@@ -139,6 +145,11 @@ class SparkPositionDeltaWrite implements DeltaWrite, RequiresDistributionAndOrde
     this.writeRequirements = writeConf.positionDeltaRequirements(command);
     this.context = new Context(dataSchema, writeConf, info, writeRequirements);
     this.writeProperties = writeConf.writeProperties();
+
+    if (this.table instanceof BaseTable) {
+      this.metricsReporter = new InMemoryMetricsReporter();
+      ((BaseTable) this.table).combineMetricsReporter(metricsReporter);
+    }
   }
 
   @Override
@@ -172,6 +183,16 @@ public DeltaBatchWrite toBatch() {
     return new PositionDeltaBatchWrite();
   }
 
+  @Override
+  public CustomMetric[] supportedCustomMetrics() {
+    return SparkWriteUtil.supportedCustomMetrics();
+  }
+
+  @Override
+  public CustomTaskMetric[] reportDriverMetrics() {
+    return SparkWriteUtil.customTaskMetrics(metricsReporter);
+  }
+
   private class PositionDeltaBatchWrite implements DeltaBatchWrite {
 
     @Override

spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkWrite.java

@@ -29,6 +29,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.BaseTable;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.FileScanTask;
@@ -53,12 +54,14 @@
 import org.apache.iceberg.io.OutputFileFactory;
 import org.apache.iceberg.io.PartitioningWriter;
 import org.apache.iceberg.io.RollingDataWriter;
+import org.apache.iceberg.metrics.InMemoryMetricsReporter;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.spark.CommitMetadata;
 import org.apache.iceberg.spark.FileRewriteCoordinator;
 import org.apache.iceberg.spark.SparkWriteConf;
 import org.apache.iceberg.spark.SparkWriteRequirements;
+import org.apache.iceberg.spark.SparkWriteUtil;
 import org.apache.iceberg.util.ContentFileUtil;
 import org.apache.iceberg.util.DataFileSet;
 import org.apache.iceberg.util.DeleteFileSet;
@@ -72,6 +75,8 @@
 import org.apache.spark.sql.catalyst.expressions.JoinedRow;
 import org.apache.spark.sql.connector.distributions.Distribution;
 import org.apache.spark.sql.connector.expressions.SortOrder;
+import org.apache.spark.sql.connector.metric.CustomMetric;
+import org.apache.spark.sql.connector.metric.CustomTaskMetric;
 import org.apache.spark.sql.connector.write.BatchWrite;
 import org.apache.spark.sql.connector.write.DataWriter;
 import org.apache.spark.sql.connector.write.DataWriterFactory;
@@ -108,6 +113,7 @@ abstract class SparkWrite implements Write, RequiresDistributionAndOrdering {
   private final Map<String, String> writeProperties;
 
   private boolean cleanupOnAbort = false;
+  private InMemoryMetricsReporter metricsReporter;
 
   SparkWrite(
       SparkSession spark,
@@ -135,6 +141,11 @@ abstract class SparkWrite implements Write, RequiresDistributionAndOrdering {
     this.writeRequirements = writeRequirements;
     this.outputSpecId = writeConf.outputSpecId();
     this.writeProperties = writeConf.writeProperties();
+
+    if (this.table instanceof BaseTable) {
+      this.metricsReporter = new InMemoryMetricsReporter();
+      ((BaseTable) this.table).combineMetricsReporter(metricsReporter);
+    }
   }
 
   @Override
@@ -163,6 +174,11 @@ public long advisoryPartitionSizeInBytes() {
     return size;
   }
 
+  @Override
+  public CustomMetric[] supportedCustomMetrics() {
+    return SparkWriteUtil.supportedCustomMetrics();
+  }
+
   BatchWrite asBatchAppend() {
     return new BatchAppend();
   }
@@ -265,6 +281,11 @@ private DataFileSet files(WriterCommitMessage[] messages) {
     return files;
   }
 
+  @Override
+  public CustomTaskMetric[] reportDriverMetrics() {
+    return SparkWriteUtil.customTaskMetrics(metricsReporter);
+  }
+
   @Override
   public String toString() {
     return String.format("IcebergWrite(table=%s, format=%s)", table, format);