viirya commented on a change in pull request #31451:
URL: https://github.com/apache/spark/pull/31451#discussion_r614246538
##########
File path:
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
##########
@@ -818,3 +818,51 @@ class ReportStatisticsDataSource extends
SimpleWritableDataSource {
}
}
}
+
+
+class SimpleCustomMetric extends CustomMetric {
+ override def name(): String = "custom_metric"
+ override def description(): String = "a simple custom metric"
+ override def aggregateTaskMetrics(taskMetrics: Array[Long]): String = {
+ s"custom_metric: ${taskMetrics.mkString(", ")}"
+ }
+}
+
+// The followings are for custom metrics of V2 data source.
+object CustomMetricReaderFactory extends PartitionReaderFactory {
+ override def createReader(partition: InputPartition):
PartitionReader[InternalRow] = {
+ val RangeInputPartition(start, end) = partition
+ new PartitionReader[InternalRow] {
+ private var current = start - 1
+
+ override def next(): Boolean = {
+ current += 1
+ current < end
+ }
+
+ override def get(): InternalRow = InternalRow(current, -current)
+
+ override def close(): Unit = {}
+
+ override def currentMetricsValues(): Array[CustomTaskMetric] = {
+ val metric = new CustomTaskMetric {
+ override def name(): String = "custom_metric"
+ override def value(): Long = 12345
+ }
+ Array(metric)
+ }
+ }
+ }
+}
+
+class CustomMetricScanBuilder extends SimpleScanBuilder {
Review comment:
Removed.
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala
##########
@@ -86,6 +92,9 @@ private class PartitionIterator[T](reader:
PartitionReader[T]) extends Iterator[
if (!hasNext) {
throw QueryExecutionErrors.endOfStreamError()
}
+ reader.currentMetricsValues.foreach { metric =>
Review comment:
yea, will do "update metrics per some rows" in a followup.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
