Spark 4.1: Add per-task write duration metrics

spark/v4.1/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRewritePositionDeleteFilesMetrics.java

@@ -0,0 +1,155 @@
+/*
+ * 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.extensions;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.iceberg.ParameterizedTestExtension;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.spark.source.SimpleRecord;
+import org.apache.iceberg.spark.source.metrics.WriteCloseDuration;
+import org.apache.iceberg.spark.source.metrics.WriteDuration;
+import org.apache.spark.sql.Encoders;
+import org.apache.spark.sql.execution.ui.SQLAppStatusStore;
+import org.apache.spark.sql.execution.ui.SQLExecutionUIData;
+import org.apache.spark.sql.execution.ui.SQLPlanMetric;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import scala.jdk.javaapi.CollectionConverters;
+
+/**
+ * Covers the write timing metrics on the position deletes rewrite path, which is driven by a
+ * procedure rather than a plain SQL write.
+ */
+@ExtendWith(ParameterizedTestExtension.class)
+public class TestRewritePositionDeleteFilesMetrics extends ExtensionsTestBase {
+
+  @AfterEach
+  public void removeTables() {
+    sql("DROP TABLE IF EXISTS %s", tableName);
+  }
+
+  @TestTemplate
+  public void writeDurationForDVRewrite() throws Exception {
+    // v3 tables rewrite deletes through DVWriter rather than DeleteWriter
+    createTableWithDeletes(3);
+
+    sql(
+        "CALL %s.system.rewrite_position_delete_files("
+            + "table => '%s', options => map('rewrite-all','true'))",
+        catalogName, tableIdent);
+
+    assertThat(maxMetricValue(new WriteDuration().description()))
+        .as("write duration should be reported by the DV rewrite")
+        .isGreaterThan(0);
+    assertThat(maxMetricValue(new WriteCloseDuration().description()))
+        .as("write close duration should be reported by the DV rewrite")
+        .isGreaterThan(0);
+  }
+
+  private void createTableWithDeletes(int formatVersion) throws Exception {
+    sql(
+        "CREATE TABLE %s (id bigint, data string) USING iceberg TBLPROPERTIES"
+            + "('format-version'='%d', 'write.delete.mode'='merge-on-read',"
+            + " 'write.delete.granularity'='partition')",
+        tableName, formatVersion);
+
+    List<SimpleRecord> records =
+        Lists.newArrayList(
+            new SimpleRecord(1, "a"),
+            new SimpleRecord(2, "b"),
+            new SimpleRecord(3, "c"),
+            new SimpleRecord(4, "d"));
+    spark
+        .createDataset(records, Encoders.bean(SimpleRecord.class))
+        .coalesce(1)
+        .writeTo(tableName)
+        .append();
+
+    sql("DELETE FROM %s WHERE id = 1", tableName);
+    sql("DELETE FROM %s WHERE id = 2", tableName);
+  }
+
+  @TestTemplate
+  public void writeDurationForPositionDeletesRewrite() throws Exception {
+    createTableWithDeletes(2);
+
+    sql(
+        "CALL %s.system.rewrite_position_delete_files("
+            + "table => '%s', options => map('rewrite-all','true'))",
+        catalogName, tableIdent);
+
+    // the rewrite runs as its own SQL execution, so look for the write node across executions
+    assertThat(maxMetricValue(new WriteDuration().description()))
+        .as("write duration should be reported by the position deletes rewrite")
+        .isGreaterThan(0);
+    assertThat(maxMetricValue(new WriteCloseDuration().description()))
+        .as("write close duration should be reported by the position deletes rewrite")
+        .isGreaterThan(0);
+  }
+
+  /** Looks up a metric by the description Spark stores in the UI, across all executions. */
+  private long maxMetricValue(String metricDescription) {
+    SQLAppStatusStore statusStore = spark.sharedState().statusStore();
+    long max = 0L;
+
+    for (SQLExecutionUIData execution : CollectionConverters.asJava(statusStore.executionsList())) {
+      Map<Object, String> metricValues =
+          CollectionConverters.asJava(
+              statusStore.execution(execution.executionId()).get().metricValues());
+
+      for (SQLPlanMetric metric : CollectionConverters.asJava(execution.metrics())) {
+        if (metric.name().equals(metricDescription)) {
+          String value = metricValues.get(metric.accumulatorId());
+          if (value != null) {
+            max = Math.max(max, parseNanos(value));
+          }
+        }
+      }
+    }
+
+    return max;
+  }
+
+  /**
+   * Durations are stored as "total (min, med, max)\n1.5 s (0.4 s, 0.5 s, 0.6 s)", so recover the
+   * total as a comparable magnitude.
+   */
+  private long parseNanos(String formatted) {
+    String total = formatted.substring(formatted.indexOf('\n') + 1);
+    total = total.substring(0, total.indexOf(" (")).trim();
+    String[] parts = total.split(" ");
+    double value = Double.parseDouble(parts[0].replace(",", ""));
+    switch (parts[1]) {
+      case "ns":
+        return (long) value;
+      case "us":
+        return (long) (value * 1_000);
+      case "ms":
+        return (long) (value * 1_000_000);
+      case "s":
+        return (long) (value * 1_000_000_000L);
+      default:
+        throw new IllegalArgumentException("Unknown duration unit: " + formatted);
+    }
+  }
+}

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

@@ -55,6 +55,8 @@
 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.spark.source.metrics.WriteCloseDuration;
+import org.apache.iceberg.spark.source.metrics.WriteDuration;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.SortOrderUtil;
 import org.apache.spark.sql.connector.distributions.Distribution;
@@ -310,12 +312,25 @@ public static CustomMetric[] supportedCustomMetrics() {
       new TotalEqualityDeletes(),
       new TotalFileSizeInBytes(),
       new TotalPositionalDeletes(),
-      new TotalRecords()
+      new TotalRecords(),
+      new WriteDuration(),
+      new WriteCloseDuration()
     };
   }
 
   public static CustomTaskMetric[] customTaskMetrics(InMemoryMetricsReporter metricsReporter) {
+    return customTaskMetrics(metricsReporter, WriteDurations.EMPTY);
+  }
+
+  /**
+   * Builds the driver-side metrics for a write, combining the commit report with the write times
+   * tasks reported back in their commit messages.
+   */
+  public static CustomTaskMetric[] customTaskMetrics(
+      InMemoryMetricsReporter metricsReporter, WriteDurations writeDurations) {
     List<CustomTaskMetric> metrics = Lists.newArrayList();
+    addValue(new WriteDuration(), writeDurations.writeNanos(), metrics);
+    addValue(new WriteCloseDuration(), writeDurations.closeNanos(), metrics);
     if (metricsReporter != null) {
       CommitReport commitReport = metricsReporter.commitReport();
       if (commitReport != null) {
@@ -348,6 +363,22 @@ public static CustomTaskMetric[] customTaskMetrics(InMemoryMetricsReporter metri
     return metrics.toArray(new CustomTaskMetric[0]);
   }
 
+  private static void addValue(
+      CustomMetric metric, long value, List<CustomTaskMetric> taskMetrics) {
+    taskMetrics.add(
+        new CustomTaskMetric() {
+          @Override
+          public String name() {
+            return metric.name();
+          }
+
+          @Override
+          public long value() {
+            return value;
+          }
+        });
+  }
+
   private static void addValue(
       CustomMetric metric, CounterResult result, List<CustomTaskMetric> taskMetrics) {
     if (result != null) {

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

@@ -0,0 +1,58 @@
+/*
+ * 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;
+
+import java.io.Serializable;
+
+/**
+ * Write times summed across the tasks of a write.
+ *
+ * <p>Tasks report their own times in their commit messages. Spark collects task metrics before it
+ * calls {@code DataWriter.commit()}, so the close time, where buffered rows are flushed, is not yet
+ * known at that point and cannot be reported as a task metric. Both times therefore travel back to
+ * the driver with the commit messages and are aggregated here.
+ */
+public class WriteDurations implements Serializable {
+
+  public static final WriteDurations EMPTY = new WriteDurations(0L, 0L);
+
+  private final long writeNanos;
+  private final long closeNanos;
+
+  private WriteDurations(long writeNanos, long closeNanos) {
+    this.writeNanos = writeNanos;
+    this.closeNanos = closeNanos;
+  }
+
+  public static WriteDurations of(long writeNanos, long closeNanos) {
+    return new WriteDurations(writeNanos, closeNanos);
+  }
+
+  public WriteDurations combine(WriteDurations other) {
+    return new WriteDurations(writeNanos + other.writeNanos, closeNanos + other.closeNanos);
+  }
+
+  public long writeNanos() {
+    return writeNanos;
+  }
+
+  public long closeNanos() {
+    return closeNanos;
+  }
+}