From f6269a8e67a2454761189f568c4c01cf467f5973 Mon Sep 17 00:00:00 2001
From: taiyang-li <taiyangli@apache.org>
Date: Tue, 21 Jul 2026 16:39:37 +0800
Subject: [PATCH] Expose shuffle reader metrics iterator delegate

---
 .../execution/ShuffledColumnarBatchRDD.scala  | 21 ++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala
index 73bc153417..ce01ce7b3e 100644
--- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala
+++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala
@@ -32,6 +32,19 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
 final private case class ShuffledColumnarBatchRDDPartition(index: Int, spec: ShufflePartitionSpec)
   extends Partition
 
+/** Wrap metrics with named class, so callers can access the delegate iterator when needed. */
+class ShuffleReaderWithMetricsIterator(
+    val delegate: Iterator[Product2[Int, ColumnarBatch]],
+    sqlMetricsReporter: SQLColumnarShuffleReadMetricsReporter)
+  extends Iterator[ColumnarBatch] {
+  override def hasNext: Boolean = delegate.hasNext
+  override def next(): ColumnarBatch = {
+    val batch = delegate.next()._2
+    sqlMetricsReporter.incBatchesRecordsRead(batch.numRows())
+    batch
+  }
+}
+
 /** [[ShuffledColumnarBatchRDD]] is the columnar version of [[org.apache.spark.rdd.ShuffledRDD]]. */
 class ShuffledColumnarBatchRDD(
     var dependency: ShuffleDependency[Int, ColumnarBatch, ColumnarBatch],
@@ -149,11 +162,9 @@ class ShuffledColumnarBatchRDD(
           sqlMetricsReporter,
           executionMode)
     }
-    reader.read().asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]].map {
-      case (_, batch: ColumnarBatch) =>
-        sqlMetricsReporter.incBatchesRecordsRead(batch.numRows())
-        batch
-    }
+    new ShuffleReaderWithMetricsIterator(
+      reader.read().asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]],
+      sqlMetricsReporter)
   }
 
   override def clearDependencies(): Unit = {