[SPARK-42510][CONNECT][PYTHON] Implement DataFrame.mapInPandas

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

@@ -60,6 +60,7 @@ message Relation {
     Unpivot unpivot = 25;
     ToSchema to_schema = 26;
     RepartitionByExpression repartition_by_expression = 27;
+    FrameMap frame_map = 28;
 
     // NA functions
     NAFill fill_na = 90;
@@ -768,3 +769,12 @@ message RepartitionByExpression {
   // (Optional) number of partitions, must be positive.
   optional int32 num_partitions = 3;
 }
+
+message FrameMap {
+  // (Required) Input relation for a Frame Map API: mapInPandas, mapInArrow.
+  Relation input = 1;
+
+  // (Required) Input user-defined function of a Frame Map API.
+  CommonInlineUserDefinedFunction func = 2;
+}
+

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

@@ -24,7 +24,7 @@ import com.google.common.collect.{Lists, Maps}
 import com.google.protobuf.{Any => ProtoAny}
 
 import org.apache.spark.TaskContext
-import org.apache.spark.api.python.SimplePythonFunction
+import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction}
 import org.apache.spark.connect.proto
 import org.apache.spark.sql.{Column, Dataset, Encoders, SparkSession}
 import org.apache.spark.sql.catalyst.{expressions, AliasIdentifier, FunctionIdentifier}
@@ -113,6 +113,8 @@ class SparkConnectPlanner(val session: SparkSession) {
       case proto.Relation.RelTypeCase.UNPIVOT => transformUnpivot(rel.getUnpivot)
       case proto.Relation.RelTypeCase.REPARTITION_BY_EXPRESSION =>
         transformRepartitionByExpression(rel.getRepartitionByExpression)
+      case proto.Relation.RelTypeCase.FRAME_MAP =>
+        transformFrameMap(rel.getFrameMap)
       case proto.Relation.RelTypeCase.RELTYPE_NOT_SET =>
         throw new IndexOutOfBoundsException("Expected Relation to be set, but is empty.")
 
@@ -465,6 +467,20 @@ class SparkConnectPlanner(val session: SparkSession) {
       .logicalPlan
   }
 
+  private def transformFrameMap(rel: proto.FrameMap): LogicalPlan = {
+    val commonUdf = rel.getFunc
+    val pythonUdf = transformPythonUDF(commonUdf)
+    pythonUdf.evalType match {
+      case PythonEvalType.SQL_MAP_PANDAS_ITER_UDF =>
+        logical.MapInPandas(
+          pythonUdf,
+          pythonUdf.dataType.asInstanceOf[StructType].toAttributes,
+          transformRelation(rel.getInput))
+      case _ =>
+        throw InvalidPlanInput(s"Function with EvalType: ${pythonUdf.evalType} is not supported")
+    }
+  }
+
   private def transformWithColumnsRenamed(rel: proto.WithColumnsRenamed): LogicalPlan = {
     Dataset
       .ofRows(session, transformRelation(rel.getInput))

dev/sparktestsupport/modules.py

@@ -531,6 +531,7 @@ def __hash__(self):
         "pyspark.sql.tests.connect.test_parity_readwriter",
         "pyspark.sql.tests.connect.test_parity_udf",
         "pyspark.sql.tests.connect.test_parity_pandas_udf",
+        "pyspark.sql.tests.connect.test_parity_pandas_map",
     ],
     excluded_python_implementations=[
         "PyPy"  # Skip these tests under PyPy since they require numpy, pandas, and pyarrow and

python/pyspark/sql/connect/_typing.py

@@ -22,10 +22,12 @@
 else:
     from typing_extensions import Protocol
 
-from typing import Any, Callable, Union, Optional
+from typing import Any, Callable, Iterable, Union, Optional
 import datetime
 import decimal
 
+from pandas.core.frame import DataFrame as PandasDataFrame
+
 from pyspark.sql.connect.column import Column
 from pyspark.sql.connect.types import DataType
 
@@ -44,6 +46,10 @@
 
 DataTypeOrString = Union[DataType, str]
 
+DataFrameLike = PandasDataFrame
+
+PandasMapIterFunction = Callable[[Iterable[DataFrameLike]], Iterable[DataFrameLike]]
+
 
 class UserDefinedFunctionLike(Protocol):
     func: Callable[..., Any]

python/pyspark/sql/connect/client.py

@@ -463,7 +463,7 @@ def register_udf(
             deterministic=deterministic,
             arguments=[],
             function=py_udf,
-        ).to_command(self)
+        ).to_plan_udf(self)
 
         # construct the request
         req = self._execute_plan_request_with_metadata()

python/pyspark/sql/connect/dataframe.py

@@ -51,6 +51,7 @@
 
 from pyspark.errors import PySparkTypeError
 from pyspark.errors.exceptions.connect import SparkConnectException
+from pyspark.rdd import PythonEvalType
 import pyspark.sql.connect.plan as plan
 from pyspark.sql.connect.group import GroupedData
 from pyspark.sql.connect.readwriter import DataFrameWriter, DataFrameWriterV2
@@ -73,6 +74,7 @@
         LiteralType,
         PrimitiveType,
         OptionalPrimitiveType,
+        PandasMapIterFunction,
     )
     from pyspark.sql.connect.session import SparkSession
 
@@ -1540,8 +1542,24 @@ def registerTempTable(self, *args: Any, **kwargs: Any) -> None:
     def storageLevel(self, *args: Any, **kwargs: Any) -> None:
         raise NotImplementedError("storageLevel() is not implemented.")
 
-    def mapInPandas(self, *args: Any, **kwargs: Any) -> None:
-        raise NotImplementedError("mapInPandas() is not implemented.")
+    def mapInPandas(
+        self, func: "PandasMapIterFunction", schema: Union[StructType, str]
+    ) -> "DataFrame":
+        from pyspark.sql.connect.udf import UserDefinedFunction
+
+        if self._plan is None:
+            raise Exception("Cannot mapInPandas when self._plan is empty.")
+
+        udf_obj = UserDefinedFunction(
+            func, returnType=schema, evalType=PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
+        )
+
+        return DataFrame.withPlan(
+            plan.FrameMap(child=self._plan, function=udf_obj, cols=self.columns),
+            session=self._session,
+        )
+
+    mapInPandas.__doc__ = PySparkDataFrame.mapInPandas.__doc__
 
     def mapInArrow(self, *args: Any, **kwargs: Any) -> None:
         raise NotImplementedError("mapInArrow() is not implemented.")

python/pyspark/sql/connect/expressions.py

@@ -547,10 +547,14 @@ def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
         )
         return expr
 
-    def to_command(self, session: "SparkConnectClient") -> "proto.CommonInlineUserDefinedFunction":
+    def to_plan_udf(self, session: "SparkConnectClient") -> "proto.CommonInlineUserDefinedFunction":
+        """Compared to `to_plan`, it returns a CommonInlineUserDefinedFunction instead of an
+        Expression."""
         expr = proto.CommonInlineUserDefinedFunction()
         expr.function_name = self._function_name
         expr.deterministic = self._deterministic
+        if len(self._arguments) > 0:
+            expr.arguments.extend([arg.to_plan(session) for arg in self._arguments])
         expr.python_udf.CopyFrom(self._function.to_plan(session))
         return expr
 

python/pyspark/sql/connect/plan.py

@@ -30,12 +30,17 @@
 
 import pyspark.sql.connect.proto as proto
 from pyspark.sql.connect.column import Column
-from pyspark.sql.connect.expressions import SortOrder, ColumnReference, LiteralExpression
+from pyspark.sql.connect.expressions import (
+    SortOrder,
+    ColumnReference,
+    LiteralExpression,
+)
 from pyspark.sql.connect.types import pyspark_types_to_proto_types
 
 if TYPE_CHECKING:
     from pyspark.sql.connect._typing import ColumnOrName
     from pyspark.sql.connect.client import SparkConnectClient
+    from pyspark.sql.connect.udf import UserDefinedFunction
 
 
 class InputValidationError(Exception):
@@ -1846,3 +1851,21 @@ def __init__(self) -> None:
 
     def plan(self, session: "SparkConnectClient") -> proto.Relation:
         return proto.Relation(catalog=proto.Catalog(list_catalogs=proto.ListCatalogs()))
+
+
+class FrameMap(LogicalPlan):
+    """Logical plan object for a Frame Map API: mapInPandas, mapInArrow."""
+
+    def __init__(
+        self, child: Optional["LogicalPlan"], function: "UserDefinedFunction", cols: List[str]
+    ) -> None:
+        super().__init__(child)
+
+        self._func = function._build_common_inline_user_defined_function(*cols)
+
+    def plan(self, session: "SparkConnectClient") -> proto.Relation:
+        assert self._child is not None
+        plan = self._create_proto_relation()
+        plan.frame_map.input.CopyFrom(self._child.plan(session))
+        plan.frame_map.func.CopyFrom(self._func.to_plan_udf(session))
+        return plan