[SPARK-37831][CORE] add task partition id in TaskInfo and Task Metrics

### Why are the changes needed?
There is no partition id in the current task metrics. It is difficult to track the task metrics of a specific partition or the stage metrics of processing data, especially when the stage was retried.
```
class TaskData private[spark](
    val taskId: Long,
    val index: Int,
    val attempt: Int,
    val launchTime: Date,
    val resultFetchStart: Option[Date],
    JsonDeserialize(contentAs = classOf[JLong])
    val duration: Option[Long],
    val executorId: String,
    val host: String,
    val status: String,
    val taskLocality: String,
    val speculative: Boolean,
    val accumulatorUpdates: Seq[AccumulableInfo],
    val errorMessage: Option[String] = None,
    val taskMetrics: Option[TaskMetrics] = None,
    val executorLogs: Map[String, String],
    val schedulerDelay: Long,
    val gettingResultTime: Long)
```

### Does this PR introduce _any_ user-facing change?
no

### How was this patch tested?
add new tests

Closes #35185 from jackylee-ch/SPARK-37831.

Lead-authored-by: stczwd <qcsd2011@163.com>
Co-authored-by: Jacky Lee <qcsd2011@gmail.com>
Co-authored-by: jackylee-ch <lijunqing@baidu.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: 2 people

core/src/main/scala/org/apache/spark/scheduler/TaskInfo.scala

@@ -35,12 +35,34 @@ class TaskInfo(
      */
     val index: Int,
     val attemptNumber: Int,
+    /**
+     * The actual RDD partition ID in this task.
+     * The ID of the RDD partition is always same across task attempts.
+     * This will be -1 for historical data, and available for all applications since Spark 3.3.
+     */
+    val partitionId: Int,
     val launchTime: Long,
     val executorId: String,
     val host: String,
     val taskLocality: TaskLocality.TaskLocality,
     val speculative: Boolean) {
 
+  /**
+   * This api doesn't contains partitionId, please use the new api.
+   * Remain it for backward compatibility before Spark 3.3.
+   */
+  def this(
+      taskId: Long,
+      index: Int,
+      attemptNumber: Int,
+      launchTime: Long,
+      executorId: String,
+      host: String,
+      taskLocality: TaskLocality.TaskLocality,
+      speculative: Boolean) = {
+    this(taskId, index, attemptNumber, -1, launchTime, executorId, host, taskLocality, speculative)
+  }
+
   /**
    * The time when the task started remotely getting the result. Will not be set if the
    * task result was sent immediately when the task finished (as opposed to sending an

core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala

@@ -509,7 +509,8 @@ private[spark] class TaskSetManager(
     // Do various bookkeeping
     copiesRunning(index) += 1
     val attemptNum = taskAttempts(index).size
-    val info = new TaskInfo(taskId, index, attemptNum, launchTime,
+    val info = new TaskInfo(
+      taskId, index, attemptNum, task.partitionId, launchTime,
       execId, host, taskLocality, speculative)
     taskInfos(taskId) = info
     taskAttempts(index) = info :: taskAttempts(index)

core/src/main/scala/org/apache/spark/status/AppStatusStore.scala

@@ -743,7 +743,8 @@ private[spark] class AppStatusStore(
       })
 
       new v1.TaskData(taskDataOld.taskId, taskDataOld.index,
-        taskDataOld.attempt, taskDataOld.launchTime, taskDataOld.resultFetchStart,
+        taskDataOld.attempt, taskDataOld.partitionId,
+        taskDataOld.launchTime, taskDataOld.resultFetchStart,
         taskDataOld.duration, taskDataOld.executorId, taskDataOld.host, taskDataOld.status,
         taskDataOld.taskLocality, taskDataOld.speculative, taskDataOld.accumulatorUpdates,
         taskDataOld.errorMessage, taskDataOld.taskMetrics,

core/src/main/scala/org/apache/spark/status/LiveEntity.scala

@@ -201,6 +201,7 @@ private class LiveTask(
       info.taskId,
       info.index,
       info.attemptNumber,
+      info.partitionId,
       info.launchTime,
       if (info.gettingResult) info.gettingResultTime else -1L,
       duration,

core/src/main/scala/org/apache/spark/status/api/v1/api.scala

@@ -308,6 +308,7 @@ class TaskData private[spark](
     val taskId: Long,
     val index: Int,
     val attempt: Int,
+    val partitionId: Int,
     val launchTime: Date,
     val resultFetchStart: Option[Date],
     @JsonDeserialize(contentAs = classOf[JLong])

core/src/main/scala/org/apache/spark/status/storeTypes.scala

@@ -141,6 +141,7 @@ private[spark] object TaskIndexNames {
   final val STAGE = "stage"
   final val STATUS = "sta"
   final val TASK_INDEX = "idx"
+  final val TASK_PARTITION_ID = "partid"
   final val COMPLETION_TIME = "ct"
 }
 
@@ -161,6 +162,10 @@ private[spark] class TaskDataWrapper(
     val index: Int,
     @KVIndexParam(value = TaskIndexNames.ATTEMPT, parent = TaskIndexNames.STAGE)
     val attempt: Int,
+    @KVIndexParam(value = TaskIndexNames.TASK_PARTITION_ID, parent = TaskIndexNames.STAGE)
+    // Different kvstores have different default values (eg 0 or -1),
+    // thus we use the default value here for backwards compatibility.
+    val partitionId: Int = -1,
     @KVIndexParam(value = TaskIndexNames.LAUNCH_TIME, parent = TaskIndexNames.STAGE)
     val launchTime: Long,
     val resultFetchStart: Long,
@@ -286,6 +291,7 @@ private[spark] class TaskDataWrapper(
       taskId,
       index,
       attempt,
+      partitionId,
       new Date(launchTime),
       if (resultFetchStart > 0L) Some(new Date(resultFetchStart)) else None,
       if (duration > 0L) Some(duration) else None,

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

@@ -317,6 +317,7 @@ private[spark] object JsonProtocol {
     ("Task ID" -> taskInfo.taskId) ~
     ("Index" -> taskInfo.index) ~
     ("Attempt" -> taskInfo.attemptNumber) ~
+    ("Partition ID" -> taskInfo.partitionId) ~
     ("Launch Time" -> taskInfo.launchTime) ~
     ("Executor ID" -> taskInfo.executorId) ~
     ("Host" -> taskInfo.host) ~
@@ -916,6 +917,7 @@ private[spark] object JsonProtocol {
     val taskId = (json \ "Task ID").extract[Long]
     val index = (json \ "Index").extract[Int]
     val attempt = jsonOption(json \ "Attempt").map(_.extract[Int]).getOrElse(1)
+    val partitionId = jsonOption(json \ "Partition ID").map(_.extract[Int]).getOrElse(-1)
     val launchTime = (json \ "Launch Time").extract[Long]
     val executorId = weakIntern((json \ "Executor ID").extract[String])
     val host = weakIntern((json \ "Host").extract[String])
@@ -930,8 +932,9 @@ private[spark] object JsonProtocol {
       case None => Seq.empty[AccumulableInfo]
     }
 
-    val taskInfo =
-      new TaskInfo(taskId, index, attempt, launchTime, executorId, host, taskLocality, speculative)
+    val taskInfo = new TaskInfo(
+      taskId, index, attempt, partitionId, launchTime,
+      executorId, host, taskLocality, speculative)
     taskInfo.gettingResultTime = gettingResultTime
     taskInfo.finishTime = finishTime
     taskInfo.failed = failed

core/src/test/resources/HistoryServerExpectations/application_list_json_expectation.json

@@ -1,4 +1,19 @@
 [ {
+  "id" : "local-1642039451826",
+  "name" : "Spark shell",
+  "attempts" : [ {
+    "startTime" : "2022-01-13T02:04:10.519GMT",
+    "endTime" : "2022-01-13T02:05:36.564GMT",
+    "lastUpdated" : "",
+    "duration" : 86045,
+    "sparkUser" : "lijunqing",
+    "completed" : true,
+    "appSparkVersion" : "3.3.0-SNAPSHOT",
+    "startTimeEpoch" : 1642039450519,
+    "endTimeEpoch" : 1642039536564,
+    "lastUpdatedEpoch" : 0
+  } ]
+}, {
   "id" : "application_1628109047826_1317105",
   "name" : "Spark shell",
   "attempts" : [ {

core/src/test/resources/HistoryServerExpectations/completed_app_list_json_expectation.json

@@ -1,4 +1,19 @@
 [ {
+  "id" : "local-1642039451826",
+  "name" : "Spark shell",
+  "attempts" : [ {
+    "startTime" : "2022-01-13T02:04:10.519GMT",
+    "endTime" : "2022-01-13T02:05:36.564GMT",
+    "lastUpdated" : "",
+    "duration" : 86045,
+    "sparkUser" : "lijunqing",
+    "completed" : true,
+    "appSparkVersion" : "3.3.0-SNAPSHOT",
+    "startTimeEpoch" : 1642039450519,
+    "endTimeEpoch" : 1642039536564,
+    "lastUpdatedEpoch" : 0
+  } ]
+}, {
   "id" : "application_1628109047826_1317105",
   "name" : "Spark shell",
   "attempts" : [ {

core/src/test/resources/HistoryServerExpectations/excludeOnFailure_for_stage_expectation.json

@@ -46,6 +46,7 @@
       "taskId" : 0,
       "index" : 0,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.347GMT",
       "duration" : 562,
       "executorId" : "0",
@@ -100,6 +101,7 @@
       "taskId" : 5,
       "index" : 3,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.958GMT",
       "duration" : 22,
       "executorId" : "1",
@@ -153,6 +155,7 @@
       "taskId" : 10,
       "index" : 8,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:19.034GMT",
       "duration" : 12,
       "executorId" : "1",
@@ -206,6 +209,7 @@
       "taskId" : 1,
       "index" : 1,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.364GMT",
       "duration" : 565,
       "executorId" : "1",
@@ -259,6 +263,7 @@
       "taskId" : 6,
       "index" : 4,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.980GMT",
       "duration" : 16,
       "executorId" : "1",
@@ -312,6 +317,7 @@
       "taskId" : 9,
       "index" : 7,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:19.022GMT",
       "duration" : 12,
       "executorId" : "1",
@@ -365,6 +371,7 @@
       "taskId" : 2,
       "index" : 2,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.899GMT",
       "duration" : 27,
       "executorId" : "0",
@@ -419,6 +426,7 @@
       "taskId" : 7,
       "index" : 5,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.996GMT",
       "duration" : 15,
       "executorId" : "1",
@@ -472,6 +480,7 @@
       "taskId" : 3,
       "index" : 0,
       "attempt" : 1,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.919GMT",
       "duration" : 24,
       "executorId" : "1",
@@ -525,6 +534,7 @@
       "taskId" : 11,
       "index" : 9,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:19.045GMT",
       "duration" : 15,
       "executorId" : "1",
@@ -578,6 +588,7 @@
       "taskId" : 8,
       "index" : 6,
       "attempt" : 0,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:19.011GMT",
       "duration" : 11,
       "executorId" : "1",
@@ -631,6 +642,7 @@
       "taskId" : 4,
       "index" : 2,
       "attempt" : 1,
+      "partitionId": -1,
       "launchTime" : "2018-01-09T10:21:18.943GMT",
       "duration" : 16,
       "executorId" : "1",