Spark 4.0: Display write metrics on SQL UI

Author: manuzhang

11340.diff

@@ -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/BaseTable.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 instanceof CommitReport) {
+      return (CommitReport) metricsReport;
+    } else {
+      return null;
+    }
+  }
 }

core/src/main/java/org/apache/iceberg/metrics/InMemoryMetricsReporter.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,101 @@ 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> taskMetrics = Lists.newArrayList();
+    if (metricsReporter != null) {
+      CommitReport commitReport = metricsReporter.commitReport();
+      if (commitReport != null) {
+        CommitMetricsResult metricsResult = commitReport.commitMetrics();
+        addMetricValue(new AddedDataFiles(), metricsResult.addedDataFiles(), taskMetrics);
+        addMetricValue(new AddedDeleteFiles(), metricsResult.addedDeleteFiles(), taskMetrics);
+        addMetricValue(
+            new AddedEqualityDeletes(), metricsResult.addedEqualityDeletes(), taskMetrics);
+        addMetricValue(
+            new AddedEqualityDeleteFiles(), metricsResult.addedEqualityDeleteFiles(), taskMetrics);
+        addMetricValue(
+            new AddedFileSizeInBytes(), metricsResult.addedFilesSizeInBytes(), taskMetrics);
+        addMetricValue(
+            new AddedPositionalDeletes(), metricsResult.addedPositionalDeletes(), taskMetrics);
+        addMetricValue(
+            new AddedPositionalDeleteFiles(),
+            metricsResult.addedPositionalDeleteFiles(),
+            taskMetrics);
+        addMetricValue(new AddedRecords(), metricsResult.addedRecords(), taskMetrics);
+        addMetricValue(new RemovedDataFiles(), metricsResult.removedDataFiles(), taskMetrics);
+        addMetricValue(new RemovedDeleteFiles(), metricsResult.removedDeleteFiles(), taskMetrics);
+        addMetricValue(new RemovedRecords(), metricsResult.removedRecords(), taskMetrics);
+        addMetricValue(
+            new RemovedEqualityDeleteFiles(),
+            metricsResult.removedEqualityDeleteFiles(),
+            taskMetrics);
+        addMetricValue(
+            new RemovedEqualityDeletes(), metricsResult.removedEqualityDeletes(), taskMetrics);
+        addMetricValue(
+            new RemovedFileSizeInBytes(), metricsResult.removedFilesSizeInBytes(), taskMetrics);
+        addMetricValue(
+            new RemovedPositionalDeleteFiles(),
+            metricsResult.removedPositionalDeleteFiles(),
+            taskMetrics);
+        addMetricValue(
+            new RemovedPositionalDeletes(), metricsResult.removedPositionalDeletes(), taskMetrics);
+        addMetricValue(new TotalDataFiles(), metricsResult.totalDataFiles(), taskMetrics);
+        addMetricValue(new TotalDeleteFiles(), metricsResult.totalDeleteFiles(), taskMetrics);
+        addMetricValue(
+            new TotalEqualityDeletes(), metricsResult.totalEqualityDeletes(), taskMetrics);
+        addMetricValue(
+            new TotalFileSizeInBytes(), metricsResult.totalFilesSizeInBytes(), taskMetrics);
+        addMetricValue(
+            new TotalPositionalDeletes(), metricsResult.totalPositionalDeletes(), taskMetrics);
+        addMetricValue(new TotalRecords(), metricsResult.totalRecords(), taskMetrics);
+      }
+    }
+    return taskMetrics.toArray(new CustomTaskMetric[0]);
+  }
+
+  private static void addMetricValue(
+      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.0/spark/src/main/java/org/apache/iceberg/spark/SparkWriteUtil.java

@@ -29,6 +29,7 @@
 import java.util.List;
 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;
@@ -65,12 +66,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;
@@ -81,6 +84,8 @@
 import org.apache.spark.sql.catalyst.InternalRow;
 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;
@@ -113,6 +118,7 @@ class SparkPositionDeltaWrite implements DeltaWrite, RequiresDistributionAndOrde
   private final Map<String, String> writeProperties;
 
   private boolean cleanupOnAbort = false;
+  private InMemoryMetricsReporter metricsReporter;
 
   SparkPositionDeltaWrite(
       SparkSession spark,
@@ -136,6 +142,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
@@ -169,6 +180,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.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWrite.java

@@ -27,6 +27,7 @@
 import java.util.Map;
 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;
@@ -51,12 +52,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.collect.ImmutableList;
 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.DataFileSet;
 import org.apache.spark.TaskContext;
 import org.apache.spark.TaskContext$;
@@ -67,6 +70,7 @@
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.connector.distributions.Distribution;
 import org.apache.spark.sql.connector.expressions.SortOrder;
+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;
@@ -103,6 +107,7 @@ abstract class SparkWrite implements Write, RequiresDistributionAndOrdering {
   private final Map<String, String> writeProperties;
 
   private boolean cleanupOnAbort = false;
+  private InMemoryMetricsReporter metricsReporter;
 
   SparkWrite(
       SparkSession spark,
@@ -130,6 +135,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
@@ -260,6 +270,11 @@ private List<DataFile> 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);

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

@@ -30,8 +30,10 @@
 import org.apache.iceberg.spark.SparkUtil;
 import org.apache.iceberg.spark.SparkWriteConf;
 import org.apache.iceberg.spark.SparkWriteRequirements;
+import org.apache.iceberg.spark.SparkWriteUtil;
 import org.apache.iceberg.types.TypeUtil;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.connector.metric.CustomMetric;
 import org.apache.spark.sql.connector.read.Scan;
 import org.apache.spark.sql.connector.write.BatchWrite;
 import org.apache.spark.sql.connector.write.LogicalWriteInfo;
@@ -159,6 +161,11 @@ public StreamingWrite toStreaming() {
           return asStreamingAppend();
         }
       }
+
+      @Override
+      public CustomMetric[] supportedCustomMetrics() {
+        return SparkWriteUtil.supportedCustomMetrics();
+      }
     };
   }
 

spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java

@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.spark.source.metrics;
+
+import org.apache.spark.sql.connector.metric.CustomSumMetric;
+
+public class AddedDataFiles extends CustomSumMetric {
+
+  static final String NAME = "addedDataFiles";
+
+  @Override
+  public String name() {
+    return NAME;
+  }
+
+  @Override
+  public String description() {
+    return "number of added data files";
+  }
+}