From 5705e1523f108e66afcf266c066615503a98a7cb Mon Sep 17 00:00:00 2001 From: Takuya UESHIN Date: Thu, 4 Jun 2020 14:53:02 -0700 Subject: [PATCH] Fix toPandas with Arrow enabled to show metrics in Query UI. --- .../src/main/scala/org/apache/spark/sql/Dataset.scala | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 12160c9f4c192..ae8d33d8558ba 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -3509,8 +3509,8 @@ class Dataset[T] private[sql]( private[sql] def collectAsArrowToR(): Array[Any] = { val timeZoneId = sparkSession.sessionState.conf.sessionLocalTimeZone - withAction("collectAsArrowToR", queryExecution) { plan => - RRDD.serveToStream("serve-Arrow") { outputStream => + RRDD.serveToStream("serve-Arrow") { outputStream => + withAction("collectAsArrowToR", queryExecution) { plan => val buffer = new ByteArrayOutputStream() val out = new DataOutputStream(outputStream) val batchWriter = new ArrowBatchStreamWriter(schema, buffer, timeZoneId) @@ -3563,8 +3563,8 @@ class Dataset[T] private[sql]( private[sql] def collectAsArrowToPython: Array[Any] = { val timeZoneId = sparkSession.sessionState.conf.sessionLocalTimeZone - withAction("collectAsArrowToPython", queryExecution) { plan => - PythonRDD.serveToStream("serve-Arrow") { outputStream => + PythonRDD.serveToStream("serve-Arrow") { outputStream => + withAction("collectAsArrowToPython", queryExecution) { plan => val out = new DataOutputStream(outputStream) val batchWriter = new ArrowBatchStreamWriter(schema, out, timeZoneId)