For review comment.

Author: viirya

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala

@@ -69,12 +69,7 @@ class DataSourceRDD(
     context.addTaskCompletionListener[Unit] { _ =>
       // In case of early stopping before consuming the entire iterator,
       // we need to do one more metric update at the end of the task.
-      reader.currentMetricsValues.foreach { metric =>
-        assert(customMetrics.contains(metric.name()),
-          s"Custom metrics ${customMetrics.keys.mkString(", ")} do not contain the metric " +
-            s"${metric.name()}")
-        customMetrics(metric.name()).set(metric.value())
-      }
+      CustomMetrics.updateMetrics(reader.currentMetricsValues, customMetrics)
       reader.close()
     }
     // TODO: SPARK-25083 remove the type erasure hack in data source scan
@@ -105,12 +100,7 @@ private class PartitionIterator[T](
       throw QueryExecutionErrors.endOfStreamError()
     }
     if (numRow % CustomMetrics.numRowsPerUpdate == 0) {
-      reader.currentMetricsValues.foreach { metric =>
-        assert(customMetrics.contains(metric.name()),
-          s"Custom metrics ${customMetrics.keys.mkString(", ")} do not contain the metric " +
-            s"${metric.name()}")
-        customMetrics(metric.name()).set(metric.value())
-      }
+      CustomMetrics.updateMetrics(reader.currentMetricsValues, customMetrics)
     }
     numRow += 1
     valuePrepared = false

sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomMetrics.scala

@@ -17,12 +17,12 @@
 
 package org.apache.spark.sql.execution.metric
 
-import org.apache.spark.sql.connector.metric.CustomMetric
+import org.apache.spark.sql.connector.metric.{CustomMetric, CustomTaskMetric}
 
 object CustomMetrics {
   private[spark] val V2_CUSTOM = "v2Custom"
 
-  private[spark] val numRowsPerUpdate = 100L
+  private[spark] val numRowsPerUpdate = 100
 
   /**
    * Given a class name, builds and returns a metric type for a V2 custom metric class
@@ -43,4 +43,15 @@ object CustomMetrics {
       None
     }
   }
+
+  /**
+   * Updates given custom metrics.
+   */
+  def updateMetrics(
+      currentMetricsValues: Seq[CustomTaskMetric],
+      customMetrics: Map[String, SQLMetric]): Unit = {
+    currentMetricsValues.foreach { metric =>
+      customMetrics(metric.name()).set(metric.value())
+    }
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousDataSourceRDD.scala

@@ -96,12 +96,7 @@ class ContinuousDataSourceRDD(
 
       override def getNext(): InternalRow = {
         if (numRow % CustomMetrics.numRowsPerUpdate == 0) {
-          partitionReader.currentMetricsValues.foreach { metric =>
-            assert(customMetrics.contains(metric.name()),
-              s"Custom metrics ${customMetrics.keys.mkString(", ")} do not contain the metric " +
-                s"${metric.name()}")
-            customMetrics(metric.name()).set(metric.value())
-          }
+          CustomMetrics.updateMetrics(partitionReader.currentMetricsValues, customMetrics)
         }
         numRow += 1
         readerForPartition.next() match {