[GLUTEN-3637][VL] Fix whole stage pipeline metric (#3661)

apache · Nov 10, 2023 · 33003d4 · 33003d4
1 parent 745ecb3
commit 33003d4
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/gluten-core/src/main/scala/io/glutenproject/utils/Iterators.scala b/gluten-core/src/main/scala/io/glutenproject/utils/Iterators.scala
@@ -20,6 +20,7 @@ import org.apache.spark.{InterruptibleIterator, TaskContext}
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.util.TaskResources
 
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
 
 private class PayloadCloser[A](in: Iterator[A])(closeCallback: A => Unit) extends Iterator[A] {
@@ -110,8 +111,10 @@ private class PipelineTimeAccumulator[A](in: Iterator[A], pipelineTime: SQLMetri
   }
 
   private def tryFinish(): Unit = {
-    pipelineTime += accumulatedTime.getAndSet(
-      0L
+    pipelineTime += TimeUnit.NANOSECONDS.toMillis(
+      accumulatedTime.getAndSet(
+        0L
+      )
     ) // make sure the accumulated time is submitted once
   }
 }