This is an automated email from the ASF dual-hosted git repository. taiyang-li pushed a commit to branch fake_add_bolt_backend in repository https://gitbox.apache.org/repos/asf/gluten.git
commit 1ae8a16664ed000e4baa5ed2db283c5fcedb01ba Author: liyang.127 <[email protected]> AuthorDate: Sat Jun 27 21:18:02 2026 +0800 [GLUTEN][CORE] Expose shuffle reader metrics iterator delegate --- .../sql/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 0642c3a247..fde17aaeed 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 @@ -28,6 +28,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], @@ -132,11 +145,9 @@ class ShuffledColumnarBatchRDD( context, sqlMetricsReporter) } - 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 = { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
