[SPARK-55619][SQL] Fix custom metrics in case of coalesced partitions

Author: peter-toth

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

@@ -34,7 +34,7 @@ class DataSourceRDDPartition(val index: Int, val inputPartition: Option[InputPar
   extends Partition with Serializable
 
 /**
- * Holds the state for a reader in a thread, used by the completion listener to access the most
+ * Holds the state for a reader in a thread/task, used by the completion listener to access the most
  * recently created reader and iterator for final metrics updates and cleanup.
  *
  * When `compute()` is called multiple times on the same thread (e.g., different input partitions
@@ -69,7 +69,7 @@ class DataSourceRDD(
     customMetrics: Map[String, SQLMetric])
   extends RDD[InternalRow](sc, Nil) {
 
-  // ThreadLocal to store the last reader state for this thread.
+  // ThreadLocal to store the last reader state for this thread/task.
   // A null value indicates that no completion listener has been added yet.
   @transient lazy private val readerStateThreadLocal = new ThreadLocal[ReaderState]()
 
@@ -85,6 +85,24 @@ class DataSourceRDD(
   }
 
   override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {
+    // In case of early stopping before consuming the entire iterator, we need to do one more metric
+    // update at the end of the task.
+    // Add completion listener only once per thread (null means no listener added yet)
+    val readerState = readerStateThreadLocal.get()
+    if (readerState == null) {
+      context.addTaskCompletionListener[Unit] { _ =>
+        // Use the reader and iterator from ThreadLocal (the last ones created in this thread/task)
+        val readerState = readerStateThreadLocal.get()
+        if (readerState != null) {
+          CustomMetrics.updateMetrics(
+            readerState.reader.currentMetricsValues.toImmutableArraySeq, customMetrics)
+          readerState.iterator.forceUpdateMetrics()
+          readerState.reader.close()
+        }
+        readerStateThreadLocal.remove()
+      }
+    }
+
     castPartition(split).inputPartition.iterator.flatMap { inputPartition =>
       val (iter, reader) = if (columnarReads) {
         val batchReader = partitionReaderFactory.createColumnarReader(inputPartition)
@@ -100,22 +118,10 @@ class DataSourceRDD(
 
       // Add completion listener only once per thread (null means no listener added yet)
       val readerState = readerStateThreadLocal.get()
-      if (readerState == null) {
-        context.addTaskCompletionListener[Unit] { _ =>
-          // Use the reader and iterator from ThreadLocal (the last ones created in this thread)
-          val readerState = readerStateThreadLocal.get()
-          if (readerState != null) {
-            // In case of early stopping before consuming the entire iterator,
-            // we need to do one more metric update at the end of the task.
-            CustomMetrics.updateMetrics(
-              readerState.reader.currentMetricsValues.toImmutableArraySeq, customMetrics)
-            readerState.iterator.forceUpdateMetrics()
-            readerState.reader.close()
-          }
-          readerStateThreadLocal.remove()
-        }
-      } else {
-        reader.initMetricsValues(readerState.reader.currentMetricsValues())
+      if (readerState != null) {
+        val metrics = readerState.reader.currentMetricsValues
+        CustomMetrics.updateMetrics(metrics.toImmutableArraySeq, customMetrics)
+        reader.initMetricsValues(metrics)
         readerState.reader.close()
       }
 

sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala

@@ -2861,6 +2861,21 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
     assert(metrics("number of rows read") == "3")
   }
 
+  test("SPARK-55619: Custom metrics of coalesced partitions") {
+    val items_partitions = Array(identity("id"))
+    createTable(items, itemsColumns, items_partitions)
+
+    sql(s"INSERT INTO testcat.ns.$items VALUES " +
+      "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+      "(2, 'bb', 10.0, cast('2021-01-01' as timestamp))")
+
+    val metrics = runAndFetchMetrics {
+      val df = sql(s"SELECT * FROM testcat.ns.$items").coalesce(1)
+      df.collect()
+    }
+    assert(metrics("number of rows read") == "2")
+  }
+
   test("SPARK-55411: Fix ArrayIndexOutOfBoundsException when join keys " +
     "are less than cluster keys") {
     withSQLConf(