Fix style

Author: mccheah

core/src/main/java/org/apache/spark/api/shuffle/ShuffleMapOutputWriter.java

@@ -20,7 +20,6 @@
 import java.io.IOException;
 
 import org.apache.spark.annotation.Experimental;
-import org.apache.spark.api.java.Optional;
 
 /**
  * :: Experimental ::

core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java

@@ -39,7 +39,6 @@
 import org.apache.spark.Partitioner;
 import org.apache.spark.ShuffleDependency;
 import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.Optional;
 import org.apache.spark.api.shuffle.SupportsTransferTo;
 import org.apache.spark.api.shuffle.ShuffleMapOutputWriter;
 import org.apache.spark.api.shuffle.ShufflePartitionWriter;

core/src/main/java/org/apache/spark/shuffle/sort/io/DefaultShuffleMapOutputWriter.java

@@ -28,15 +28,13 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.Optional;
 import org.apache.spark.api.shuffle.ShuffleMapOutputWriter;
 import org.apache.spark.api.shuffle.ShufflePartitionWriter;
 import org.apache.spark.api.shuffle.SupportsTransferTo;
 import org.apache.spark.api.shuffle.TransferrableWritableByteChannel;
 import org.apache.spark.internal.config.package$;
 import org.apache.spark.shuffle.sort.DefaultTransferrableWritableByteChannel;
 import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
-import org.apache.spark.storage.BlockManagerId;
 import org.apache.spark.shuffle.IndexShuffleBlockResolver;
 import org.apache.spark.storage.TimeTrackingOutputStream;
 import org.apache.spark.util.Utils;

core/src/main/java/org/apache/spark/shuffle/sort/io/DefaultShuffleWriteSupport.java

@@ -22,7 +22,6 @@
 import org.apache.spark.api.shuffle.ShuffleMapOutputWriter;
 import org.apache.spark.api.shuffle.ShuffleWriteSupport;
 import org.apache.spark.shuffle.IndexShuffleBlockResolver;
-import org.apache.spark.storage.BlockManagerId;
 
 public class DefaultShuffleWriteSupport implements ShuffleWriteSupport {
 
@@ -41,8 +40,13 @@ public ShuffleMapOutputWriter createMapOutputWriter(
       int shuffleId,
       int mapId,
       int numPartitions) {
+    TaskContext taskContext = TaskContext.get();
+    if (taskContext == null) {
+      throw new IllegalStateException(
+          "Task context must be set before creating a map output writer.");
+    }
     return new DefaultShuffleMapOutputWriter(
       shuffleId, mapId, numPartitions,
-        TaskContext.get().taskMetrics().shuffleWriteMetrics(), blockResolver, sparkConf);
+        taskContext.taskMetrics().shuffleWriteMetrics(), blockResolver, sparkConf);
   }
 }