Review feedback

Author: sryza

core/src/main/scala/org/apache/spark/executor/TaskMetrics.scala

@@ -163,6 +163,16 @@ object DataReadMethod extends Enumeration with Serializable {
   val Memory, Disk, Hadoop, Network = Value
 }
 
+/**
+ * :: DeveloperApi ::
+ * Method by which output data was written.
+ */
+@DeveloperApi
+object DataWriteMethod extends Enumeration with Serializable {
+  type DataWriteMethod = Value
+  val Hadoop = Value
+}
+
 /**
  * :: DeveloperApi ::
  * Metrics about reading input data.
@@ -180,7 +190,7 @@ case class InputMetrics(readMethod: DataReadMethod.Value) {
  * Metrics about writing output data.
  */
 @DeveloperApi
-case class OutputMetrics() {
+case class OutputMetrics(writeMethod: DataWriteMethod.Value) {
   /**
    * Total bytes written
    */

core/src/main/scala/org/apache/spark/rdd/PairRDDFunctions.scala

@@ -40,7 +40,7 @@ import org.apache.spark.Partitioner.defaultPartitioner
 import org.apache.spark.SparkContext._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.deploy.SparkHadoopUtil
-import org.apache.spark.executor.OutputMetrics
+import org.apache.spark.executor.{DataWriteMethod, OutputMetrics}
 import org.apache.spark.partial.{BoundedDouble, PartialResult}
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.util.Utils
@@ -978,13 +978,7 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
       val committer = format.getOutputCommitter(hadoopContext)
       committer.setupTask(hadoopContext)
 
-      val bytesWrittenCallback = Option(config.get("mapreduce.output.fileoutputformat.outputdir"))
-        .map(new Path(_))
-        .flatMap(SparkHadoopUtil.get.getFSBytesWrittenOnThreadCallback(_, config))
-      val outputMetrics = new OutputMetrics()
-      if (bytesWrittenCallback.isDefined) {
-        context.taskMetrics.outputMetrics = Some(outputMetrics)
-      }
+      val (outputMetrics, bytesWrittenCallback) = initHadoopOutputMetrics(context, config)
 
       val writer = format.getRecordWriter(hadoopContext).asInstanceOf[NewRecordWriter[K,V]]
       try {
@@ -1061,13 +1055,7 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
       // around by taking a mod. We expect that no task will be attempted 2 billion times.
       val attemptNumber = (context.attemptId % Int.MaxValue).toInt
 
-      val bytesWrittenCallback = Option(config.get("mapreduce.output.fileoutputformat.outputdir"))
-        .map(new Path(_))
-        .flatMap(SparkHadoopUtil.get.getFSBytesWrittenOnThreadCallback(_, config))
-      val outputMetrics = new OutputMetrics()
-      if (bytesWrittenCallback.isDefined) {
-        context.taskMetrics.outputMetrics = Some(outputMetrics)
-      }
+      val (outputMetrics, bytesWrittenCallback) = initHadoopOutputMetrics(context, config)
 
       writer.setup(context.stageId, context.partitionId, attemptNumber)
       writer.open()
@@ -1098,6 +1086,32 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
     writer.commitJob()
   }
 
+  private def initHadoopOutputMetrics(context: TaskContext, config: Configuration)
+    : (OutputMetrics, Option[() => Long]) = {
+    val bytesWrittenCallback = Option(config.get("mapreduce.output.fileoutputformat.outputdir"))
+      .map(new Path(_))
+      .flatMap(SparkHadoopUtil.get.getFSBytesWrittenOnThreadCallback(_, config))
+    val outputMetrics = new OutputMetrics(DataWriteMethod.Hadoop)
+    if (bytesWrittenCallback.isDefined) {
+      context.taskMetrics.outputMetrics = Some(outputMetrics)
+    }
+    (outputMetrics, bytesWrittenCallback)
+  }
+
+  /*
+  private def maybeUpdateOutputMetrics(recordsWritten: Long) {
+    // Update bytes written metric every few records
+    if (recordsSinceMetricsUpdate ==
+      PairRDDFunctions.RECORDS_BETWEEN_BYTES_WRITTEN_METRIC_UPDATES
+      && bytesWrittenCallback.isDefined) {
+      recordsSinceMetricsUpdate = 0
+      bytesWrittenCallback.foreach { fn => outputMetrics.bytesWritten = fn() }
+    } else {
+      recordsSinceMetricsUpdate += 1
+    }
+
+  }*/
+
   /**
    * Return an RDD with the keys of each tuple.
    */

core/src/main/scala/org/apache/spark/util/JsonProtocol.scala

@@ -272,6 +272,7 @@ private[spark] object JsonProtocol {
   }
 
   def outputMetricsToJson(outputMetrics: OutputMetrics): JValue = {
+    ("Data Write Method" -> outputMetrics.writeMethod.toString) ~
     ("Bytes Written" -> outputMetrics.bytesWritten)
   }
 
@@ -621,7 +622,8 @@ private[spark] object JsonProtocol {
   }
 
   def outputMetricsFromJson(json: JValue): OutputMetrics = {
-    val metrics = new OutputMetrics()
+    val metrics = new OutputMetrics(
+      DataWriteMethod.withName((json \ "Data Write Method").extract[String]))
     metrics.bytesWritten = (json \ "Bytes Written").extract[Long]
     metrics
   }

core/src/test/scala/org/apache/spark/ui/jobs/JobProgressListenerSuite.scala

@@ -159,7 +159,7 @@ class JobProgressListenerSuite extends FunSuite with LocalSparkContext with Matc
       val inputMetrics = new InputMetrics(DataReadMethod.Hadoop)
       taskMetrics.inputMetrics = Some(inputMetrics)
       inputMetrics.bytesRead = base + 7
-      val outputMetrics = new OutputMetrics()
+      val outputMetrics = new OutputMetrics(DataWriteMethod.Hadoop)
       taskMetrics.outputMetrics = Some(outputMetrics)
       outputMetrics.bytesWritten = base + 8
       taskMetrics

core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala

@@ -617,7 +617,7 @@ class JsonProtocolSuite extends FunSuite {
       t.setShuffleReadMetrics(Some(sr))
     }
     if (hasOutput) {
-      val outputMetrics = new OutputMetrics()
+      val outputMetrics = new OutputMetrics(DataWriteMethod.Hadoop)
       outputMetrics.bytesWritten = a + b + c
       t.outputMetrics = Some(outputMetrics)
     } else {
@@ -1022,6 +1022,7 @@ class JsonProtocolSuite extends FunSuite {
       |      "Bytes Read": 2100
       |    },
       |    "Output Metrics": {
+      |      "Data Write Method": "Hadoop",
       |      "Bytes Written": 1200
       |    },
       |    "Updated Blocks": [