[SPARK-46812][CONNECT][PYTHON] Make mapInPandas / mapInArrow support ResourceProfile

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

@@ -967,6 +967,34 @@ message FetchErrorDetailsResponse {
   }
 }
 
+message BuildResourceProfileRequest {
+  // (Required)
+  //
+  // The session_id specifies a spark session for a user id (which is specified
+  // by user_context.user_id). The session_id is set by the client to be able to
+  // collate streaming responses from different queries within the dedicated session.
+  // The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff`
+  string session_id = 1;
+
+  // (Required) User context
+  UserContext user_context = 2;
+
+  // (Required) The ResourceProfile to be built on the server-side.
+  ResourceProfile profile = 3;
+}
+
+// Response to building resource profile.
+message BuildResourceProfileResponse {
+  string session_id = 1;
+
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 4;
+
+  // (Required) Server-side generated resource profile id.
+  int32 profile_id = 5;
+}
+
 // Main interface for the SparkConnect service.
 service SparkConnectService {
 
@@ -1011,5 +1039,7 @@ service SparkConnectService {
 
   // FetchErrorDetails retrieves the matched exception with details based on a provided error id.
   rpc FetchErrorDetails(FetchErrorDetailsRequest) returns (FetchErrorDetailsResponse) {}
-}
 
+  // Build ResourceProfile and get the profile id
+  rpc BuildResourceProfile(BuildResourceProfileRequest) returns (BuildResourceProfileResponse) {}
+}

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

@@ -46,3 +46,38 @@ message ResourceInformation {
   // (Required) An array of strings describing the addresses of the resource.
   repeated string addresses = 2;
 }
+
+// An executor resource request.
+message ExecutorResourceRequest {
+  // (Required) resource name.
+  string resource_name = 1;
+
+  // (Required) resource amount requesting.
+  int64 amount = 2;
+
+  // Optional script used to discover the resources.
+  optional string discovery_script = 3;
+
+  // Optional vendor, required for some cluster managers.
+  optional string vendor = 4;
+}
+
+// A task resource request.
+message TaskResourceRequest {
+  // (Required) resource name.
+  string resource_name = 1;
+
+  // (Required) resource amount requesting as a double to support fractional
+  // resource requests.
+  double amount = 2;
+}
+
+message ResourceProfile {
+  // (Optional) Resource requests for executors. Mapped from the resource name
+  // (e.g., cores, memory, CPU) to its specific request.
+  map<string, ExecutorResourceRequest> executor_resources = 1;
+
+  // (Optional) Resource requests for tasks. Mapped from the resource name
+  // (e.g., cores, memory, CPU) to its specific request.
+  map<string, TaskResourceRequest> task_resources = 2;
+}

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

@@ -23,6 +23,7 @@ import "google/protobuf/any.proto";
 import "spark/connect/expressions.proto";
 import "spark/connect/types.proto";
 import "spark/connect/catalog.proto";
+import "spark/connect/common.proto";
 
 option java_multiple_files = true;
 option java_package = "org.apache.spark.connect.proto";
@@ -892,6 +893,9 @@ message MapPartitions {
 
   // (Optional) Whether to use barrier mode execution or not.
   optional bool is_barrier = 3;
+
+  // (Optional) ResourceProfile id used for the stage level scheduling.
+  optional int32 profile_id = 4;
 }
 
 message GroupMap {

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

@@ -543,21 +543,27 @@ class SparkConnectPlanner(
       case proto.CommonInlineUserDefinedFunction.FunctionCase.PYTHON_UDF =>
         val pythonUdf = transformPythonUDF(commonUdf)
         val isBarrier = if (rel.hasIsBarrier) rel.getIsBarrier else false
+        val profile = if (rel.hasProfileId) {
+          val profileId = rel.getProfileId
+          Some(session.sparkContext.resourceProfileManager.resourceProfileFromId(profileId))
+        } else {
+          None
+        }
         pythonUdf.evalType match {
           case PythonEvalType.SQL_MAP_PANDAS_ITER_UDF =>
             logical.MapInPandas(
               pythonUdf,
               DataTypeUtils.toAttributes(pythonUdf.dataType.asInstanceOf[StructType]),
               baseRel,
               isBarrier,
-              None)
+              profile)
           case PythonEvalType.SQL_MAP_ARROW_ITER_UDF =>
             logical.MapInArrow(
               pythonUdf,
               DataTypeUtils.toAttributes(pythonUdf.dataType.asInstanceOf[StructType]),
               baseRel,
               isBarrier,
-              None)
+              profile)
           case _ =>
             throw InvalidPlanInput(
               s"Function with EvalType: ${pythonUdf.evalType} is not supported")

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectBuildResourceProfileHandler.scala

@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connect.service
+
+import scala.jdk.CollectionConverters.MapHasAsScala
+
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.resource.{ExecutorResourceRequest, ResourceProfile, TaskResourceProfile, TaskResourceRequest}
+
+class SparkConnectBuildResourceProfileHandler(
+    responseObserver: StreamObserver[proto.BuildResourceProfileResponse])
+    extends Logging {
+
+  /**
+   * transform the spark connect ResourceProfile to spark ResourceProfile
+   * @param rp
+   *   Spark connect ResourceProfile
+   * @return
+   *   the Spark ResourceProfile
+   */
+  private def transformResourceProfile(rp: proto.ResourceProfile): ResourceProfile = {
+    val ereqs = rp.getExecutorResourcesMap.asScala.map { case (name, res) =>
+      name -> new ExecutorResourceRequest(
+        res.getResourceName,
+        res.getAmount,
+        res.getDiscoveryScript,
+        res.getVendor)
+    }.toMap
+    val treqs = rp.getTaskResourcesMap.asScala.map { case (name, res) =>
+      name -> new TaskResourceRequest(res.getResourceName, res.getAmount)
+    }.toMap
+
+    if (ereqs.isEmpty) {
+      new TaskResourceProfile(treqs)
+    } else {
+      new ResourceProfile(ereqs, treqs)
+    }
+  }
+
+  def handle(request: proto.BuildResourceProfileRequest): Unit = {
+    val holder = SparkConnectService
+      .getOrCreateIsolatedSession(request.getUserContext.getUserId, request.getSessionId)
+
+    val rp = transformResourceProfile(request.getProfile)
+
+    val session = holder.session
+    session.sparkContext.resourceProfileManager.addResourceProfile(rp)
+
+    val builder = proto.BuildResourceProfileResponse.newBuilder()
+    builder.setProfileId(rp.id)
+    builder.setSessionId(request.getSessionId)
+    builder.setServerSideSessionId(holder.serverSessionId)
+    responseObserver.onNext(builder.build())
+    responseObserver.onCompleted()
+  }
+
+}

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

@@ -33,7 +33,7 @@ import org.apache.commons.lang3.StringUtils
 
 import org.apache.spark.{SparkContext, SparkEnv}
 import org.apache.spark.connect.proto
-import org.apache.spark.connect.proto.{AddArtifactsRequest, AddArtifactsResponse, SparkConnectServiceGrpc}
+import org.apache.spark.connect.proto.{AddArtifactsRequest, AddArtifactsResponse, BuildResourceProfileRequest, BuildResourceProfileResponse, SparkConnectServiceGrpc}
 import org.apache.spark.connect.proto.SparkConnectServiceGrpc.AsyncService
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config.UI.UI_ENABLED
@@ -227,6 +227,20 @@ class SparkConnectService(debug: Boolean) extends AsyncService with BindableServ
     }
   }
 
+  override def buildResourceProfile(
+      request: BuildResourceProfileRequest,
+      responseObserver: StreamObserver[BuildResourceProfileResponse]): Unit = {
+    try {
+      new SparkConnectBuildResourceProfileHandler(responseObserver).handle(request)
+    } catch {
+      ErrorUtils.handleError(
+        "buildResourceProfile",
+        observer = responseObserver,
+        userId = request.getUserContext.getUserId,
+        sessionId = request.getSessionId)
+    }
+  }
+
   private def methodWithCustomMarshallers(methodDesc: MethodDescriptor[MessageLite, MessageLite])
       : MethodDescriptor[MessageLite, MessageLite] = {
     val recursionLimit =

dev/check_pyspark_custom_errors.py

@@ -176,6 +176,7 @@ def check_pyspark_custom_errors(target_paths, exclude_paths):
     TARGET_PATHS = ["python/pyspark/sql"]
     EXCLUDE_PATHS = [
         "python/pyspark/sql/tests",
+        "python/pyspark/sql/connect/resource",
         "python/pyspark/sql/connect/proto",
     ]
 

dev/sparktestsupport/modules.py

@@ -554,6 +554,7 @@ def __hash__(self):
         "pyspark.resource.profile",
         # unittests
         "pyspark.resource.tests.test_resources",
+        "pyspark.resource.tests.test_connect_resources",
     ],
 )
 
@@ -1027,6 +1028,7 @@ def __hash__(self):
         "pyspark.sql.tests.connect.test_parity_pandas_udf_scalar",
         "pyspark.sql.tests.connect.test_parity_pandas_udf_grouped_agg",
         "pyspark.sql.tests.connect.test_parity_pandas_udf_window",
+        "pyspark.sql.tests.connect.test_resources",
     ],
     excluded_python_implementations=[
         "PyPy"  # Skip these tests under PyPy since they require numpy, pandas, and pyarrow and

python/pyspark/resource/profile.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-
+from threading import RLock
 from typing import overload, Dict, Union, Optional
 
 from py4j.java_gateway import JavaObject
@@ -37,6 +37,9 @@ class ResourceProfile:
 
     .. versionadded:: 3.1.0
 
+    .. versionchanged:: 4.0.0
+        Supports Spark Connect.
+
     Notes
     -----
     This API is evolving.
@@ -99,6 +102,11 @@ def __init__(
         _exec_req: Optional[Dict[str, ExecutorResourceRequest]] = None,
         _task_req: Optional[Dict[str, TaskResourceRequest]] = None,
     ):
+        # profile id
+        self._id: Optional[int] = None
+        # lock to protect _id
+        self._lock = RLock()
+
         if _java_resource_profile is not None:
             self._java_resource_profile = _java_resource_profile
         else:
@@ -114,14 +122,23 @@ def id(self) -> int:
         int
             A unique id of this :class:`ResourceProfile`
         """
+        with self._lock:
+            if self._id is None:
+                if self._java_resource_profile is not None:
+                    self._id = self._java_resource_profile.id()
+                else:
+                    from pyspark.sql import is_remote
 
-        if self._java_resource_profile is not None:
-            return self._java_resource_profile.id()
-        else:
-            raise RuntimeError(
-                "SparkContext must be created to get the id, get the id "
-                "after adding the ResourceProfile to an RDD"
-            )
+                    if is_remote():
+                        from pyspark.sql.connect.resource.profile import ResourceProfile
+
+                        rp = ResourceProfile(
+                            self._executor_resource_requests, self._task_resource_requests
+                        )
+                        self._id = rp.id
+                    else:
+                        raise RuntimeError("SparkContext must be created to get the profile id.")
+            return self._id
 
     @property
     def taskResources(self) -> Dict[str, TaskResourceRequest]:
@@ -185,7 +202,10 @@ def __init__(self) -> None:
 
         # TODO: ignore[attr-defined] will be removed, once SparkContext is inlined
         _jvm = SparkContext._jvm
-        if _jvm is not None:
+
+        from pyspark.sql import is_remote
+
+        if _jvm is not None and not is_remote():
             self._jvm = _jvm
             self._java_resource_profile_builder = (
                 _jvm.org.apache.spark.resource.ResourceProfileBuilder()

python/pyspark/resource/requests.py

@@ -164,9 +164,11 @@ def __init__(
         _requests: Optional[Dict[str, ExecutorResourceRequest]] = None,
     ):
         from pyspark import SparkContext
+        from pyspark.sql import is_remote
 
         _jvm = _jvm or SparkContext._jvm
-        if _jvm is not None:
+
+        if _jvm is not None and not is_remote():
             self._java_executor_resource_requests = (
                 _jvm.org.apache.spark.resource.ExecutorResourceRequests()
             )
@@ -460,9 +462,11 @@ def __init__(
         _requests: Optional[Dict[str, TaskResourceRequest]] = None,
     ):
         from pyspark import SparkContext
+        from pyspark.sql import is_remote
 
         _jvm = _jvm or SparkContext._jvm
-        if _jvm is not None:
+
+        if _jvm is not None and not is_remote():
             self._java_task_resource_requests: Optional[
                 JavaObject
             ] = _jvm.org.apache.spark.resource.TaskResourceRequests()