From b8c1abe5f55caae081911d2c427b14d6b93e944c Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Tue, 27 Oct 2026 19:23:39 +0800 Subject: [PATCH] [spark] Fix read batch time for the first batch and active splits --- .../paimon/spark/PaimonPartitionReader.scala | 22 +++-- .../spark/PaimonRecordReaderIterator.scala | 3 +- .../paimon/spark/sql/PaimonMetricTest.scala | 93 ++++++++++++++++++- 3 files changed, 108 insertions(+), 10 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala index 02ac56f57888..7f8b6ca05ba4 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala @@ -104,8 +104,7 @@ case class PaimonPartitionReader( if (currentRow != null) { stop = true } else { - totalReadBatchTimeMs += currentRecordReader.readBatchTimeMs - currentRecordReader.close() + closeCurrentReader() currentRecordReader = readSplit() if (currentRecordReader == null) { stop = true @@ -135,15 +134,24 @@ case class PaimonPartitionReader( } override def currentMetricsValues(): Array[CustomTaskMetric] = { - partitionMetrics ++ Array(PaimonReadBatchTimeTaskMetric(totalReadBatchTimeMs)) + val currentReadBatchTimeMs = + if (currentRecordReader == null) 0L else currentRecordReader.readBatchTimeMs + partitionMetrics ++ Array( + PaimonReadBatchTimeTaskMetric(totalReadBatchTimeMs + currentReadBatchTimeMs)) + } + + private def closeCurrentReader(): Unit = { + if (currentRecordReader != null) { + val reader = currentRecordReader + reader.close() + totalReadBatchTimeMs += reader.readBatchTimeMs + currentRecordReader = null + } } override def close(): Unit = { try { - if (currentRecordReader != null) { - totalReadBatchTimeMs += currentRecordReader.readBatchTimeMs - currentRecordReader.close() - } + closeCurrentReader() } finally { ioManager.close() } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala index 1c6dc813159e..b31d48ba3b6d 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala @@ -78,10 +78,11 @@ case class PaimonRecordReaderIterator( private val joinedRow: JoinedRow = JoinedRow.join(null, metadataRow) private var lastFilePath: Path = _ + // Initialize before readBatch() so the first batch's elapsed time is not reset to zero. + private var readBatchTimeNs: Long = 0L private var currentIterator: RecordReader.RecordIterator[PaimonInternalRow] = readBatch() private var advanced = false private var currentResult: PaimonInternalRow = _ - private var readBatchTimeNs: Long = 0L override def hasNext: Boolean = { if (currentIterator == null) { diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonMetricTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonMetricTest.scala index ede81d7ffb7e..bb23a3c4145a 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonMetricTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonMetricTest.scala @@ -18,12 +18,17 @@ package org.apache.paimon.spark.sql +import org.apache.paimon.data.{InternalRow => PaimonRow} +import org.apache.paimon.disk.IOManager +import org.apache.paimon.reader.RecordReader +import org.apache.paimon.reader.RecordReader.RecordIterator +import org.apache.paimon.spark.{PaimonInputPartition, PaimonPartitionReader, PaimonRecordReaderIterator, PaimonSparkTestBase} import org.apache.paimon.spark.PaimonMetrics.{RESULTED_TABLE_FILES, SCANNED_SNAPSHOT_ID, SKIPPED_TABLE_FILES} -import org.apache.paimon.spark.PaimonSparkTestBase import org.apache.paimon.spark.metric.SparkMetricRegistry import org.apache.paimon.spark.read.PaimonSplitScan import org.apache.paimon.spark.util.ScanPlanHelper -import org.apache.paimon.table.source.DataSplit +import org.apache.paimon.table.source.{DataSplit, ReadBuilder, Split, TableRead} +import org.apache.paimon.types.{DataTypes, RowType} import org.apache.spark.metrics.source.PaimonMetricsSource import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} @@ -35,13 +40,97 @@ import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.paimon.Utils import org.junit.jupiter.api.Assertions +import org.mockito.ArgumentMatchers.any +import org.mockito.Mockito.{doThrow, mock, times, verify, when} import javax.management.ObjectName +import java.io.IOException import java.lang.management.ManagementFactory class PaimonMetricTest extends PaimonSparkTestBase with ScanPlanHelper { + test("read batch time includes the first batch") { + val reader = timedReader() + val iterator = PaimonRecordReaderIterator(reader, Seq.empty, mock(classOf[Split])) + try { + assert(iterator.readBatchTimeMs >= 20L) + } finally { + iterator.close() + } + } + + test("read batch time includes active splits without counting them again on close") { + for (earlyClose <- Seq(true, false); failClose <- Seq(false, true)) { + val inputs = Seq(timedReader(), timedReader()) + if (failClose) { + doThrow(new IOException("close failed")) + .doNothing() + .when(inputs.last) + .close() + } + val splits = inputs.map(_ => mock(classOf[Split])) + val read = mock(classOf[TableRead]) + when(read.withIOManager(any(classOf[IOManager]))).thenReturn(read) + splits.zip(inputs).foreach { + case (split, input) => when(read.createReader(split)).thenReturn(input) + } + val builder = mock(classOf[ReadBuilder]) + when(builder.readType()).thenReturn(RowType.of(DataTypes.INT())) + when(builder.newRead()).thenReturn(read) + val reader = PaimonPartitionReader( + builder, + PaimonInputPartition(splits), + Seq.empty, + true, + null, + Array.empty) + def readTime: Long = metric(reader.currentMetricsValues(), "readBatchTime") + try { + assert(reader.next() && reader.get().getInt(0) == 1) + val first = readTime + assert(first >= 20L) + assert(readTime == first) + assert(reader.next() && reader.get().getInt(0) == 2) + assert(reader.next() && reader.get().getInt(0) == 1) + assert(readTime >= first + 20L) + if (!earlyClose) { + assert(reader.next() && reader.get().getInt(0) == 2) + if (failClose) intercept[IOException](reader.next()) + else assert(!reader.next()) + } + val beforeClose = readTime + assert(beforeClose >= first + 20L) + if (earlyClose && failClose) { + intercept[IOException](reader.close()) + assert(readTime == beforeClose) + } + reader.close() + assert(readTime == beforeClose) + reader.close() + assert(readTime == beforeClose) + verify(inputs.head).close() + verify(inputs.last, times(if (failClose) 2 else 1)).close() + } finally { + reader.close() + } + } + } + + private def timedReader(): RecordReader[PaimonRow] = { + val batch = mock(classOf[RecordIterator[PaimonRow]]) + when(batch.next()).thenReturn(GenericRow.of(1), GenericRow.of(2), null) + val reader = mock(classOf[RecordReader[PaimonRow]]) + when(reader.readBatch()) + .thenAnswer { + _ => + Thread.sleep(25) + batch + } + .thenReturn(null) + reader + } + test("Paimon Metric: registered metrics are exposed through JMX") { val tableName = "jmx-metric-test-" + System.nanoTime() val registry = SparkMetricRegistry()