Unify the behavior of pyspark.TaskContext with spark core

In Spark Core, there is a `TaskContext` object which is a singleton. We set a task context instance which can be TaskContext or BarrierTaskContext before the task function startup, and unset it to none after the function end. So we can both get TaskContext and BarrierTaskContext with the object. How we can only get the BarrierTaskContext with `BarrierTaskContext`, we will get `None` if we get it by `TaskContext.get` in a barrier stage.

In this patch, we unify the behavior of TaskContext for pyspark with Spark core. This is useful when people switch from normal code to barrier code, and only need a little update.