[SPARK-24822][PySpark] Python support for barrier execution mode

core/src/main/scala/org/apache/spark/rdd/RDDBarrier.scala

@@ -19,9 +19,9 @@ package org.apache.spark.rdd
 
 import scala.reflect.ClassTag
 
-import org.apache.spark.BarrierTaskContext
 import org.apache.spark.TaskContext
 import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.api.java.JavaRDD
 
 /** Represents an RDD barrier, which forces Spark to launch tasks of this stage together. */
 class RDDBarrier[T: ClassTag](rdd: RDD[T]) {
@@ -47,5 +47,18 @@ class RDDBarrier[T: ClassTag](rdd: RDD[T]) {
     )
   }
 
+  /**
+   * Expose a JavaRDD that wraps a barrier RDD generated from the prev RDD, to support launch
+   * barrier stage from python side.
+   */
+  private[spark] def toJavaRDD(): JavaRDD[T] = {
+    val barrierRDD = new MapPartitionsRDD[T, T](
+      rdd,
+      (context, pid, iter) => iter,
+      preservesPartitioning = false,
+      isFromBarrier = true)
+    JavaRDD.fromRDD(barrierRDD)
+  }
+
   /** TODO extra conf(e.g. timeout) */
 }

python/pyspark/rdd.py

@@ -2406,6 +2406,26 @@ def toLocalIterator(self):
             sock_info = self.ctx._jvm.PythonRDD.toLocalIteratorAndServe(self._jrdd.rdd())
         return _load_from_socket(sock_info, self._jrdd_deserializer)
 
+    def barrier(self):
+        """
+        .. note:: Experimental
+
+        Indicates that Spark must launch the tasks together for the current stage.
+
+        .. versionadded:: 2.4.0
+        """
+        return RDDBarrier(self)
+
+    def isBarrier(self):
+        """
+        .. note:: Experimental
+
+        Whether this RDD is in a barrier stage.
+
+        .. versionadded:: 2.4.0
+        """
+        return self._jrdd.rdd().isBarrier()
+
 
 def _prepare_for_python_RDD(sc, command):
     # the serialized command will be compressed by broadcast
@@ -2429,6 +2449,36 @@ def _wrap_function(sc, func, deserializer, serializer, profiler=None):
                                   sc.pythonVer, broadcast_vars, sc._javaAccumulator)
 
 
+class RDDBarrier(object):
+
+    """
+    .. note:: Experimental
+
+    An RDDBarrier turns an RDD into a barrier RDD, which forces Spark to launch tasks of the stage
+    contains this RDD together.
+
+    .. versionadded:: 2.4.0
+    """
+
+    def __init__(self, rdd):
+        self.rdd = rdd
+        self._jrdd = rdd._jrdd
+
+    def mapPartitions(self, f, preservesPartitioning=False):
+        """
+        .. note:: Experimental
+
+        Return a new RDD by applying a function to each partition of this RDD.
+
+        .. versionadded:: 2.4.0
+        """
+        def func(s, iterator):
+            return f(iterator)
+        jBarrierRdd = self._jrdd.rdd().barrier().toJavaRDD()
+        pyBarrierRdd = RDD(jBarrierRdd, self.rdd.ctx, self.rdd._jrdd_deserializer)
+        return pyBarrierRdd.mapPartitions(f, preservesPartitioning)
+
+
 class PipelinedRDD(RDD):
 
     """