Fix confusing getChunks().head

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

@@ -81,9 +81,8 @@ private[spark] class TaskResultGetter(sparkEnv: SparkEnv, scheduler: TaskSchedul
                   taskSetManager, tid, TaskState.FINISHED, TaskResultLost)
                 return
               }
-              // TODO(josh): assumption that there is only one chunk here is a hack
               val deserializedResult = serializer.get().deserialize[DirectTaskResult[_]](
-                serializedTaskResult.get.getChunks().head)
+                serializedTaskResult.get.toByteBuffer)
               sparkEnv.blockManager.master.removeBlock(blockId)
               (deserializedResult, size)
           }

core/src/main/scala/org/apache/spark/util/io/ChunkedByteBuffer.scala

@@ -108,7 +108,7 @@ private[spark] class ChunkedByteBuffer(var chunks: Array[ByteBuffer]) {
   }
 
   /**
-   * Get duplicates of the ByteBuffers backing this ChunkedByteBuffer. Visible for testing
+   * Get duplicates of the ByteBuffers backing this ChunkedByteBuffer.
    */
   def getChunks(): Array[ByteBuffer] = {
     chunks.map(_.duplicate())