[SPARK-25921][PySpark] Fix barrier task run without BarrierTaskContext while python worker reuse

python/pyspark/taskcontext.py

@@ -144,10 +144,19 @@ def __init__(self):
         """Construct a BarrierTaskContext, use get instead"""
         pass
 
+    def __new__(cls):
+        """
+        Rewrite __new__ method to BarrierTaskContext for _getOrCreate called when _taskContext
+        is not instance of BarrierTaskContext.
+        """
+        if not isinstance(cls._taskContext, BarrierTaskContext):
+            cls._taskContext = object.__new__(cls)
+        return cls._taskContext
+
     @classmethod
     def _getOrCreate(cls):
         """Internal function to get or create global BarrierTaskContext."""
-        if cls._taskContext is None:
+        if not isinstance(cls._taskContext, BarrierTaskContext):
             cls._taskContext = BarrierTaskContext()
         return cls._taskContext
 

python/pyspark/tests.py

@@ -614,6 +614,18 @@ def context_barrier(x):
         times = rdd.barrier().mapPartitions(f).map(context_barrier).collect()
         self.assertTrue(max(times) - min(times) < 1)
 
+    def test_barrier_with_python_worker_reuse(self):
+        """
+        Verify that BarrierTaskContext.barrier() with reused python worker.
+        """
+        rdd = self.sc.parallelize(range(4), 4)
+        # start a normal job first to start all worker
+        result = rdd.map(lambda x: x ** 2).collect()
+        self.assertEqual([0, 1, 4, 9], result)
+
+        # worker will be reused in this barrier job
+        self.test_barrier()
+
     def test_barrier_infos(self):
         """
         Verify that BarrierTaskContext.getTaskInfos() returns a list of all task infos in the