[SPARK-3543] Clean up Java TaskContext implementation.

This addresses some minor issues in https://github.com/apache/spark/pull/2425

Author: Reynold Xin <rxin@apache.org>

Closes #2557 from rxin/TaskContext and squashes the following commits:

a51e5f6 [Reynold Xin] [SPARK-3543] Clean up Java TaskContext implementation.

1 files changed, 14 insertions, 19 deletions

diff --git a/core/src/main/java/org/apache/spark/TaskContext.java b/core/src/main/java/org/apache/spark/TaskContext.java
index 09b8ce02bd..4e6d708af0 100644
--- a/core/src/main/java/org/apache/spark/TaskContext.java
+++ b/core/src/main/java/org/apache/spark/TaskContext.java
@@ -56,7 +56,7 @@ public class TaskContext implements Serializable {
    * @param taskMetrics performance metrics of the task
    */
   @DeveloperApi
-  public TaskContext(Integer stageId, Integer partitionId, Long attemptId, Boolean runningLocally,
+  public TaskContext(int stageId, int partitionId, long attemptId, boolean runningLocally,
                      TaskMetrics taskMetrics) {
     this.attemptId = attemptId;
     this.partitionId = partitionId;
@@ -65,7 +65,6 @@ public class TaskContext implements Serializable {
     this.taskMetrics = taskMetrics;
   }
 
-
   /**
    * :: DeveloperApi ::
    * Contextual information about a task which can be read or mutated during execution.
@@ -76,8 +75,7 @@ public class TaskContext implements Serializable {
    * @param runningLocally whether the task is running locally in the driver JVM
    */
   @DeveloperApi
-  public TaskContext(Integer stageId, Integer partitionId, Long attemptId,
-                     Boolean runningLocally) {
+  public TaskContext(int stageId, int partitionId, long attemptId, boolean runningLocally) {
     this.attemptId = attemptId;
     this.partitionId = partitionId;
     this.runningLocally = runningLocally;
@@ -85,7 +83,6 @@ public class TaskContext implements Serializable {
     this.taskMetrics = TaskMetrics.empty();
   }
 
-
   /**
    * :: DeveloperApi ::
    * Contextual information about a task which can be read or mutated during execution.
@@ -95,7 +92,7 @@ public class TaskContext implements Serializable {
    * @param attemptId the number of attempts to execute this task
    */
   @DeveloperApi
-  public TaskContext(Integer stageId, Integer partitionId, Long attemptId) {
+  public TaskContext(int stageId, int partitionId, long attemptId) {
     this.attemptId = attemptId;
     this.partitionId = partitionId;
     this.runningLocally = false;
@@ -107,9 +104,9 @@ public class TaskContext implements Serializable {
     new ThreadLocal<TaskContext>();
 
   /**
-  * :: Internal API ::
-  * This is spark internal API, not intended to be called from user programs.
-  */
+   * :: Internal API ::
+   * This is spark internal API, not intended to be called from user programs.
+   */
   public static void setTaskContext(TaskContext tc) {
     taskContext.set(tc);
   }
@@ -118,10 +115,8 @@ public class TaskContext implements Serializable {
     return taskContext.get();
   }
 
-  /** 
-  * :: Internal API ::
-  */
-  public static void remove() {
+  /** :: Internal API ::  */
+  public static void unset() {
     taskContext.remove();
   }
 
@@ -130,22 +125,22 @@ public class TaskContext implements Serializable {
     new ArrayList<TaskCompletionListener>();
 
   // Whether the corresponding task has been killed.
-  private volatile Boolean interrupted = false;
+  private volatile boolean interrupted = false;
 
   // Whether the task has completed.
-  private volatile Boolean completed = false;
+  private volatile boolean completed = false;
 
   /**
    * Checks whether the task has completed.
    */
-  public Boolean isCompleted() {
+  public boolean isCompleted() {
     return completed;
   }
 
   /**
    * Checks whether the task has been killed.
    */
-  public Boolean isInterrupted() {
+  public boolean isInterrupted() {
     return interrupted;
   }
 
@@ -246,12 +241,12 @@ public class TaskContext implements Serializable {
   }
 
   @Deprecated
-  /** Deprecated: use getRunningLocally() */
+  /** Deprecated: use isRunningLocally() */
   public boolean runningLocally() {
     return runningLocally;
   }
 
-  public boolean getRunningLocally() {
+  public boolean isRunningLocally() {
     return runningLocally;
   }