[SPARK-56598][SQL] Custom metrics support for TruncatableTable

### What changes were proposed in this pull request?

Added support for connectors to provide custom metrics in Truncate and Delete operations. Refactored the existing custom metrics logic to be reusable across different plan nodes.

### Why are the changes needed?

For connectors to provide metric values.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Existing tests.

### Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Opus 4.7

Closes #55511 from ZiyaZa/custom-metrics.

Authored-by: Ziya Mukhtarov <ziya5muxtarov@gmail.com>
Signed-off-by: Anton Okolnychyi <aokolnychyi@apache.org>

Author: ZiyaZa

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TruncatableTable.java

@@ -18,6 +18,8 @@
 package org.apache.spark.sql.connector.catalog;
 
 import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.metric.CustomMetric;
+import org.apache.spark.sql.connector.metric.CustomTaskMetric;
 
 /**
  * Represents a table which can be atomically truncated.
@@ -34,4 +36,25 @@ public interface TruncatableTable extends Table {
    * @since 3.2.0
    */
   boolean truncateTable();
+
+  /**
+   * Returns an array of supported custom metrics with name and description.
+   * By default it returns empty array.
+   *
+   * @since 4.2.0
+   */
+  default CustomMetric[] supportedCustomMetrics() {
+    return new CustomMetric[]{};
+  }
+
+  /**
+   * Returns an array of custom metrics which are collected with values at the driver side only.
+   * Note that these metrics must be included in the supported custom metrics reported by
+   * `supportedCustomMetrics`.
+   *
+   * @since 4.2.0
+   */
+  default CustomTaskMetric[] reportDriverMetrics() {
+    return new CustomTaskMetric[]{};
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala

@@ -129,7 +129,7 @@ case class BatchScanExec(
       new DataSourceRDD(
         sparkContext, filteredPartitions, readerFactory, supportsColumnar, customMetrics)
     }
-    postDriverMetrics()
+    postDriverMetrics(scan.reportDriverMetrics())
     rdd
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ContinuousScanExec.scala

@@ -65,7 +65,7 @@ case class ContinuousScanExec(
       schema,
       readerFactory,
       customMetrics)
-    postDriverMetrics()
+    postDriverMetrics(scan.reportDriverMetrics())
     inputRDD
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala

@@ -24,23 +24,22 @@ import org.apache.spark.sql.catalyst.plans.physical
 import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, PartitionReaderFactory, Scan}
-import org.apache.spark.sql.execution.{ExplainUtils, LeafExecNode, SafeForKWayMerge, SQLExecution}
-import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.execution.{ExplainUtils, LeafExecNode, SafeForKWayMerge}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.internal.connector.SupportsMetadata
 import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.spark.util.ArrayImplicits._
 import org.apache.spark.util.Utils
 
-trait DataSourceV2ScanExecBase extends LeafExecNode with SafeForKWayMerge {
+trait DataSourceV2ScanExecBase
+  extends LeafExecNode
+  with SafeForKWayMerge
+  with SupportsCustomDriverMetrics {
 
-  lazy val customMetrics = scan.supportedCustomMetrics().map { customMetric =>
-    customMetric.name() -> SQLMetrics.createV2CustomMetric(sparkContext, customMetric)
-  }.toMap
+  override lazy val customMetrics: Map[String, SQLMetric] =
+    createCustomMetrics(scan.supportedCustomMetrics())
 
-  override lazy val metrics = {
-    Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) ++
-      customMetrics
-  }
+  override protected lazy val sparkMetrics: Map[String, SQLMetric] =
+    Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))
 
   def scan: Scan
 
@@ -145,18 +144,6 @@ trait DataSourceV2ScanExecBase extends LeafExecNode with SafeForKWayMerge {
     }
   }
 
-  protected def postDriverMetrics(): Unit = {
-    val driveSQLMetrics = scan.reportDriverMetrics().map(customTaskMetric => {
-      val metric = metrics(customTaskMetric.name())
-      metric.set(customTaskMetric.value())
-      metric
-    })
-
-    val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
-    SQLMetrics.postDriverMetricUpdates(sparkContext, executionId,
-      driveSQLMetrics.toImmutableArraySeq)
-  }
-
   override def doExecuteColumnar(): RDD[ColumnarBatch] = {
     val numOutputRows = longMetric("numOutputRows")
     inputRDD.asInstanceOf[RDD[ColumnarBatch]].map { b =>

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DeleteFromTableExec.scala

@@ -23,18 +23,29 @@ import org.apache.spark.sql.catalyst.transactions.TransactionUtils
 import org.apache.spark.sql.connector.catalog.SupportsDeleteV2
 import org.apache.spark.sql.connector.catalog.transactions.Transaction
 import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.execution.metric.SQLMetric
 
 case class DeleteFromTableExec(
     table: SupportsDeleteV2,
     condition: Array[Predicate],
     refreshCache: () => Unit,
-    transaction: Option[Transaction] = None) extends LeafV2CommandExec with TransactionalExec {
+    transaction: Option[Transaction] = None)
+  extends LeafV2CommandExec
+  with TransactionalExec
+  with SupportsCustomDriverMetrics {
+
+  override lazy val customMetrics: Map[String, SQLMetric] =
+    createCustomMetrics(table.supportedCustomMetrics())
 
   override def withTransaction(txn: Option[Transaction]): DeleteFromTableExec =
     copy(transaction = txn)
 
   override protected def run(): Seq[InternalRow] = {
-    table.deleteWhere(condition)
+    try {
+      table.deleteWhere(condition)
+    } finally {
+      postDriverMetrics(table.reportDriverMetrics())
+    }
     transaction.foreach(TransactionUtils.commit)
     refreshCache()
     Seq.empty

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/MicroBatchScanExec.scala

@@ -55,7 +55,7 @@ case class MicroBatchScanExec(
   override lazy val inputRDD: RDD[InternalRow] = {
     val inputRDD = new DataSourceRDD(sparkContext, partitions, readerFactory, supportsColumnar,
       customMetrics)
-    postDriverMetrics()
+    postDriverMetrics(scan.reportDriverMetrics())
     inputRDD
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RealTimeStreamScanExec.scala

@@ -166,7 +166,7 @@ case class RealTimeStreamScanExec(
       supportsColumnar,
       customMetrics
     )
-    postDriverMetrics()
+    postDriverMetrics(scan.reportDriverMetrics())
     inputRDD
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/SupportsCustomDriverMetrics.scala

@@ -0,0 +1,80 @@
+/*
+ * 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.spark.sql.execution.datasources.v2
+
+import org.apache.spark.sql.connector.metric.{CustomMetric, CustomTaskMetric}
+import org.apache.spark.sql.execution.{SparkPlan, SQLExecution}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.util.ArrayImplicits._
+
+/**
+ * A mixin for Spark plan nodes that expose driver-side custom metrics reported by a connector.
+ * Implementations declare the connector-owned metrics via [[customMetrics]]; after the underlying
+ * operation has executed they call [[postDriverMetrics]] with the connector's reported values so
+ * they are visible in the SQL UI.
+ *
+ * Nodes that also expose Spark-owned metrics supply them via [[sparkMetrics]]. Names in
+ * [[sparkMetrics]] are reserved: if the connector happens to report a value under the same name,
+ * Spark's value wins and the connector's is dropped.
+ */
+trait SupportsCustomDriverMetrics { self: SparkPlan =>
+
+  /**
+   * The custom metrics the connector supports for this operation, keyed by name.
+   */
+  def customMetrics: Map[String, SQLMetric]
+
+  /**
+   * Spark-owned metrics that should appear alongside the connector-declared ones. Values under
+   * these names are owned by Spark and take precedence on a name collision.
+   */
+  protected def sparkMetrics: Map[String, SQLMetric] = Map.empty
+
+  override lazy val metrics: Map[String, SQLMetric] = customMetrics ++ sparkMetrics
+
+  /**
+   * Converts an array of connector-declared metrics into the map shape [[customMetrics]] uses.
+   */
+  protected def createCustomMetrics(metrics: Array[CustomMetric]): Map[String, SQLMetric] = {
+    metrics.map { m =>
+      m.name -> SQLMetrics.createV2CustomMetric(sparkContext, m)
+    }.toMap
+  }
+
+  /**
+   * Applies the values reported by the connector to the declared metrics and posts them so the
+   * SQL UI reflects the final values. Metrics not declared via [[customMetrics]] are ignored.
+   * Metrics whose name collides with [[sparkMetrics]] are also ignored so Spark-owned values
+   * are preserved.
+   */
+  protected def postDriverMetrics(taskMetrics: Array[CustomTaskMetric]): Unit = {
+    val updated = taskMetrics.flatMap { t =>
+      if (sparkMetrics.contains(t.name())) {
+        // Spark metrics take precedence on collisions.
+        None
+      } else {
+        metrics.get(t.name()).map { metric =>
+          metric.set(t.value())
+          metric
+        }
+      }
+    }
+    val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+    SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, updated.toImmutableArraySeq)
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/TruncateTableExec.scala

@@ -20,18 +20,28 @@ package org.apache.spark.sql.execution.datasources.v2
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.connector.catalog.TruncatableTable
+import org.apache.spark.sql.execution.metric.SQLMetric
 
 /**
  * Physical plan node for table truncation.
  */
 case class TruncateTableExec(
     table: TruncatableTable,
-    refreshCache: () => Unit) extends LeafV2CommandExec {
+    refreshCache: () => Unit)
+  extends LeafV2CommandExec
+  with SupportsCustomDriverMetrics {
+
+  override lazy val customMetrics: Map[String, SQLMetric] =
+    createCustomMetrics(table.supportedCustomMetrics())
 
   override def output: Seq[Attribute] = Seq.empty
 
   override protected def run(): Seq[InternalRow] = {
-    if (table.truncateTable()) refreshCache()
-    Seq.empty
+    try {
+      if (table.truncateTable()) refreshCache()
+      Seq.empty
+    } finally {
+      postDriverMetrics(table.reportDriverMetrics())
+    }
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V1FallbackWriters.scala

@@ -23,8 +23,8 @@ import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.classic.Dataset
 import org.apache.spark.sql.connector.catalog.SupportsWrite
 import org.apache.spark.sql.connector.write.V1Write
-import org.apache.spark.sql.execution.{SparkPlan, SQLExecution}
-import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.sources.InsertableRelation
 
 /**
@@ -56,13 +56,15 @@ case class OverwriteByExpressionExecV1(
     write: V1Write) extends V1FallbackWriters
 
 /** Some helper interfaces that use V2 write semantics through the V1 writer interface. */
-sealed trait V1FallbackWriters extends LeafV2CommandExec with SupportsV1Write {
+sealed trait V1FallbackWriters
+  extends LeafV2CommandExec
+  with SupportsV1Write
+  with SupportsCustomDriverMetrics {
+
   override def output: Seq[Attribute] = Nil
 
-  override val metrics: Map[String, SQLMetric] =
-    write.supportedCustomMetrics().map { customMetric =>
-      customMetric.name() -> SQLMetrics.createV2CustomMetric(sparkContext, customMetric)
-    }.toMap
+  override lazy val customMetrics: Map[String, SQLMetric] =
+    createCustomMetrics(write.supportedCustomMetrics())
 
   def table: SupportsWrite
   def refreshCache: () => Unit
@@ -75,12 +77,7 @@ sealed trait V1FallbackWriters extends LeafV2CommandExec with SupportsV1Write {
 
       Nil
     } finally {
-      write.reportDriverMetrics().foreach { customTaskMetric =>
-        metrics.get(customTaskMetric.name()).foreach(_.set(customTaskMetric.value()))
-      }
-
-      val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
-      SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq)
+      postDriverMetrics(write.reportDriverMetrics())
     }
   }
 }