[SPARK-42929][CONNECT] make mapInPandas / mapInArrow support "is_barrier"

connector/connect/common/src/main/protobuf/spark/connect/relations.proto

@@ -794,6 +794,9 @@ message MapPartitions {
 
   // (Required) Input user-defined function.
   CommonInlineUserDefinedFunction func = 2;
+
+  // (Optional) isBarrier.
+  optional bool is_barrier = 3;
 }
 
 message GroupMap {

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -484,19 +484,20 @@ class SparkConnectPlanner(val session: SparkSession) {
   private def transformMapPartitions(rel: proto.MapPartitions): LogicalPlan = {
     val commonUdf = rel.getFunc
     val pythonUdf = transformPythonUDF(commonUdf)
+    val isBarrier = if (rel.hasIsBarrier) rel.getIsBarrier else false
     pythonUdf.evalType match {
       case PythonEvalType.SQL_MAP_PANDAS_ITER_UDF =>
         logical.MapInPandas(
           pythonUdf,
           pythonUdf.dataType.asInstanceOf[StructType].toAttributes,
           transformRelation(rel.getInput),
-          false)
+          isBarrier)
       case PythonEvalType.SQL_MAP_ARROW_ITER_UDF =>
         logical.PythonMapInArrow(
           pythonUdf,
           pythonUdf.dataType.asInstanceOf[StructType].toAttributes,
           transformRelation(rel.getInput),
-          false)
+          isBarrier)
       case _ =>
         throw InvalidPlanInput(s"Function with EvalType: ${pythonUdf.evalType} is not supported")
     }

python/pyspark/sql/connect/dataframe.py

@@ -1623,6 +1623,7 @@ def _map_partitions(
         func: "PandasMapIterFunction",
         schema: Union[StructType, str],
         evalType: int,
+        is_barrier: bool,
     ) -> "DataFrame":
         from pyspark.sql.connect.udf import UserDefinedFunction
 
@@ -1636,21 +1637,31 @@ def _map_partitions(
         )
 
         return DataFrame.withPlan(
-            plan.MapPartitions(child=self._plan, function=udf_obj, cols=self.columns),
+            plan.MapPartitions(
+                child=self._plan, function=udf_obj, cols=self.columns, is_barrier=is_barrier
+            ),
             session=self._session,
         )
 
     def mapInPandas(
-        self, func: "PandasMapIterFunction", schema: Union[StructType, str]
+        self,
+        func: "PandasMapIterFunction",
+        schema: Union[StructType, str],
+        is_barrier: bool = False,
     ) -> "DataFrame":
-        return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF)
+        return self._map_partitions(
+            func, schema, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, is_barrier
+        )
 
     mapInPandas.__doc__ = PySparkDataFrame.mapInPandas.__doc__
 
     def mapInArrow(
-        self, func: "ArrowMapIterFunction", schema: Union[StructType, str]
+        self,
+        func: "ArrowMapIterFunction",
+        schema: Union[StructType, str],
+        is_barrier: bool = False,
     ) -> "DataFrame":
-        return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_ARROW_ITER_UDF)
+        return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_ARROW_ITER_UDF, is_barrier)
 
     mapInArrow.__doc__ = PySparkDataFrame.mapInArrow.__doc__
 

python/pyspark/sql/connect/plan.py

@@ -1912,17 +1912,23 @@ class MapPartitions(LogicalPlan):
     """Logical plan object for a mapPartitions-equivalent API: mapInPandas, mapInArrow."""
 
     def __init__(
-        self, child: Optional["LogicalPlan"], function: "UserDefinedFunction", cols: List[str]
+        self,
+        child: Optional["LogicalPlan"],
+        function: "UserDefinedFunction",
+        cols: List[str],
+        is_barrier: bool,
     ) -> None:
         super().__init__(child)
 
         self._func = function._build_common_inline_user_defined_function(*cols)
+        self._is_barrier = is_barrier
 
     def plan(self, session: "SparkConnectClient") -> proto.Relation:
         assert self._child is not None
         plan = self._create_proto_relation()
         plan.map_partitions.input.CopyFrom(self._child.plan(session))
         plan.map_partitions.func.CopyFrom(self._func.to_plan_udf(session))
+        plan.map_partitions.is_barrier = self._is_barrier
         return plan