This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 51a9e96516 [GLUTEN][CORE] Expose shuffle reader metrics iterator
delegate (#12432)
51a9e96516 is described below
commit 51a9e96516d4d7414a52027d2019acf46e2b6009
Author: 李扬 <[email protected]>
AuthorDate: Fri Jul 3 22:51:33 2026 +0800
[GLUTEN][CORE] Expose shuffle reader metrics iterator delegate (#12432)
Change-Id: Ia535725b902e511da1693bd9f7fbaa38f84cc3c8
Co-authored-by: Aime <[email protected]>
---
.../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]
