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 36beabd411 [GLUTEN-12462][CORE] Expose shuffle reader metrics iterator
delegate (#12587)
36beabd411 is described below
commit 36beabd4118501e686f1c5ae700ce6b18a0e1b3b
Author: 李扬 <[email protected]>
AuthorDate: Tue Jul 21 21:14:50 2026 +0800
[GLUTEN-12462][CORE] Expose shuffle reader metrics iterator delegate
(#12587)
---
.../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 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 = {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
