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

[wbo4958]

What changes were proposed in this pull request?

Support stage-level scheduling for some PySpark DataFrame APIs (mapInPandas and mapInArrow).

Why are the changes needed?

The introduction of barrier mode in Spark, as seen in #40520, allows for the implementation of Spark ML cases (pure Python algorithms) using DataFrame APIs such as mapInPandas and mapInArrow, so it's necessary to enable stage-level scheduling for DataFrame APIs.

Does this PR introduce any user-facing change?

Yes, This PR adds a new argument "profile" for mapInPandas and mapInArrow.

def mapInPandas(
    self, func: "PandasMapIterFunction",
        schema: Union[StructType, str],
        barrier: bool = False,
        profile: Optional[ResourceProfile] = None,
) -> "DataFrame":

def mapInArrow(
    self, func: "ArrowMapIterFunction",
        schema: Union[StructType, str],
        barrier: bool = False,
        profile: Optional[ResourceProfile] = None,
) -> "DataFrame":

How to use it? take mapInPandas as an example,

from pyspark import TaskContext
def func(iterator):
    tc = TaskContext.get()
    assert tc.cpus() == 3
    for batch in iterator:
        yield batch
df = spark.range(10)

from pyspark.resource import TaskResourceRequests, ResourceProfileBuilder
treqs = TaskResourceRequests().cpus(3)
rp = ResourceProfileBuilder().require(treqs).build

df.mapInPandas(func, "id long", False, rp).collect()

How was this patch tested?

The newly added tests can pass, and some manual tests are needed for dynamic allocation on or off.

Was this patch authored or co-authored using generative AI tooling?

No

[wbo4958]

I will perform the manual tests for this PR, will update it after done.

[wbo4958]

I was trying to add unit tests to check if the ResourceProfile is correctly applied to the underlying RDD generated in MapInPandasExec, here is my testing code,

        df1 = df.mapInPandas(lambda iter: iter, "id long")
        assert df1.rdd.getResourceProfile() is None

        treqs = TaskResourceRequests().cpus(2)
        expected_rp = ResourceProfileBuilder().require(treqs).build

        df2 = df.mapInPandas(lambda iter: iter, "id long", False, expected_rp)
        assert df2.rdd.getResourceProfile() is not None

But the ResourceProfile got from df2.rdd.getResourceProfile() is None, the reason for it is df2.rdd will add some other extra MapPartitionRDDs that don't have ResourceProfile attached.

I also tried to use JVM RDD to get the correct parent RDD with the below code,

df2.rdd._jrdd.firstParent()

or

df2.rdd._jrdd.parent(0)

But both of them didn't work, with below error messages,

py4j.protocol.Py4JError: An error occurred while calling o45.parent. Trace:
py4j.Py4JException: Method parent([class java.lang.Integer]) does not exist

y4j.protocol.Py4JError: An error occurred while calling o45.firstParent. Trace:
py4j.Py4JException: Method firstParent([]) does not exist

I don't know how to add unit tests for this PR, but I will perform the manual tests.

[wbo4958]

Manual tests

The manual tests were conducted on a spark Standalone cluster with only 1 worker.

With dynamic allocation disabled.

pyspark --master spark://192.168.141.19:7077 --conf spark.executor.cores=4 --conf spark.task.cpus=1 \
   --conf spark.dynamicAllocation.enabled=false

The above command requires 1 executor with 4 CPU cores, and the default task.cpus = 1, so the default tasks parallelism is 4 at a time.

1. task.cores=1

Test code:

from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder

def filter_func(iterator):
    for pdf in iterator:
        yield pdf

df = spark.range(0, 100, 1, 6)
from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder
treqs = TaskResourceRequests().cpus(1)
rp = ResourceProfileBuilder().require(treqs).build
df.repartition(3).mapInArrow(lambda iter: iter, df.schema, False, rp).collect()

When the required task.cpus=1, executor.cores=4 (No executor resource specified, use the default one), there will be 4 tasks running for rp at the same time.

The entire Spark application consists of a single Spark job that will be divided into two stages. The first shuffle stage comprises 6 tasks, the first 4 tasks will be executed simultaneously, then the last 2 tasks.

The second ResultStage comprises 3 tasks, all of which will be executed simultaneously since the required task.cpus is 1.

2. task.cores=2

Test code,

from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder

def filter_func(iterator):
    for pdf in iterator:
        yield pdf

df = spark.range(0, 100, 1, 6)
from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder
treqs = TaskResourceRequests().cpus(2)
rp = ResourceProfileBuilder().require(treqs).build
df.repartition(3).mapInArrow(lambda iter: iter, df.schema, False, rp).collect()

When the required task.cpus=2, executor.cores=4 (No executor resource specified, use the default one), there will be 2 tasks running for rp.

The first shuffle stage behaves the same as the first one.

The second ResultStage comprises 3 tasks, so the first 2 tasks will be running at a time, and then execute the last task.

3. task.cores=3

Test code,

from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder

def filter_func(iterator):
    for pdf in iterator:
        yield pdf

df = spark.range(0, 100, 1, 6)
from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder
treqs = TaskResourceRequests().cpus(3)
rp = ResourceProfileBuilder().require(treqs).build
df.repartition(3).mapInArrow(lambda iter: iter, df.schema, False, rp).collect()

When the required task.cpus=3, executor.cores=4 (No executor resource specified, use the default one), there will be 1 task running for rp.

The first shuffle stage behaves the same as the first one.

The second ResultStage comprises 3 tasks, all of which will be running serially.

4. task.cores=5

from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder

def filter_func(iterator):
    for pdf in iterator:
        yield pdf

df = spark.range(0, 100, 1, 6)
from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder
treqs = TaskResourceRequests().cpus(5)
rp = ResourceProfileBuilder().require(treqs).build
df.repartition(3).mapInArrow(lambda iter: iter, df.schema, False, rp).collect()

exception happened.

4j.protocol.Py4JJavaError: An error occurred while calling o160.collectToPython.
: org.apache.spark.SparkException: The number of cores per executor (=4) has to be >= the number of cpus per task = 5.
	at org.apache.spark.resource.ResourceUtils$.validateTaskCpusLargeEnough(ResourceUtils.scala:413)
	at org.apache.spark.resource.ResourceProfile.calculateTasksAndLimitingResource(ResourceProfile.scala:193)
	at org.apache.spark.resource.ResourceProfile.$anonfun$limitingResource$1(ResourceProfile.scala:163)
	at scala.Option.getOrElse(Option.scala:201)
	at org.apache.spark.resource.ResourceProfile.limitingResource(ResourceProfile.scala:162)
	at org.apache.spark.resource.ResourceProfileManager.addResourceProfile(ResourceProfileManager.scala:142)
	at org.apache.spark.rdd.RDD.withResources(RDD.scala:1844)
	at org.apache.spark.sql.execution.python.MapInBatchExec.$anonfun$doExecute$3(MapInBatchExec.scala:86)
	at scala.Option.map(Option.scala:242)
	at org.apache.spark.sql.execution.python.MapInBatchExec.doExecute(MapInBatchExec.scala:86)
	at org.apache.spark.sql.execution.python.MapInBatchExec.doExecute$(MapInBatchExec.scala:50)
	at org.apache.spark.sql.execution.python.MapInArrowExec.doExecute(MapInArrowExec.scala:29)
	at org.apache.spark.sql.execution.SparkPlan.$anonfun$execute$1(SparkPlan.scala:195)
	at org.apache.spark.sql.execution.SparkPlan.$anonfun$executeQuery$1(SparkPlan.scala:246)
	at org.apache.spark.rdd.RDDOperationScope$.withScope(RDDOperationScope.scala:151)
	at org.apache.spark.sql.execution.SparkPlan.executeQuery(SparkPlan.scala:243)
	at org.apache.spark.sql.execution.SparkPlan.execute(SparkPlan.scala:191)
	at org.apache.spark.sql.execution.SparkPlan.getByteArrayRdd(SparkPlan.scala:364)
	at org.apache.spark.sql.execution.SparkPlan.executeCollect(SparkPlan.scala:445)
	at org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.$anonfun$executeCollect$1(AdaptiveSparkPlanExec.scala:390)
	at org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.withFinalPlanUpdate(AdaptiveSparkPlanExec.scala:418)
	at org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.executeCollect(AdaptiveSparkPlanExec.scala:390)
	at org.apache.spark.sql.Dataset.$anonfun$collectToPython$1(Dataset.scala:4265)
	at org.apache.spark.sql.Dataset.$anonfun$withAction$2(Dataset.scala:4439)
	at org.apache.spark.sql.execution.QueryExecution$.withInternalError(QueryExecution.scala:560)
	at org.apache.spark.sql.Dataset.$anonfun$withAction$1(Dataset.scala:4437)
	at org.apache.spark.sql.execution.SQLExecution$.$anonfun$withNewExecutionId0$6(SQLExecution.scala:150)
	at org.apache.spark.sql.execution.SQLExecution$.withSQLConfPropagated(SQLExecution.scala:241)
	at org.apache.spark.sql.execution.SQLExecution$.$anonfun$withNewExecutionId0$1(SQLExecution.scala:116)
	at org.apache.spark.sql.SparkSession.withActive(SparkSession.scala:918)
	at org.apache.spark.sql.execution.SQLExecution$.withNewExecutionId0(SQLExecution.scala:72)
	at org.apache.spark.sql.execution.SQLExecution$.withNewExecutionId(SQLExecution.scala:196)
	at org.apache.spark.sql.Dataset.withAction(Dataset.scala:4437)
	at org.apache.spark.sql.Dataset.collectToPython(Dataset.scala:4262)
	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
	at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:77)
	at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
	at java.base/java.lang.reflect.Method.invoke(Method.java:568)
	at py4j.reflection.MethodInvoker.invoke(MethodInvoker.java:244)
	at py4j.reflection.ReflectionEngine.invoke(ReflectionEngine.java:374)
	at py4j.Gateway.invoke(Gateway.java:282)
	at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132)
	at py4j.commands.CallCommand.execute(CallCommand.java:79)
	at py4j.ClientServerConnection.waitForCommands(ClientServerConnection.java:182)
	at py4j.ClientServerConnection.run(ClientServerConnection.java:106)
	at java.base/java.lang.Thread.run(Thread.java:840)

[wbo4958]

With dynamic allocation enabled.

pyspark --master spark://192.168.141.19:7077 --conf spark.executor.cores=4 --conf spark.task.cpus=1 \
  --conf spark.dynamicAllocation.enabled=true --conf spark.dynamicAllocation.maxExecutors=1\

The above command enables the dynamic allocation and the max executors required is set to 1 in order to test.

TaskResourceProfile without any specific executor request information

Test code,

from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder

def filter_func(iterator):
    for pdf in iterator:
        yield pdf

df = spark.range(0, 100, 1, 4)

treqs = TaskResourceRequests().cpus(3)
rp = ResourceProfileBuilder().require(treqs).build
df.repartition(3).mapInArrow(lambda iter: iter, df.schema, False, rp).collect()

The rp refers to the TaskResourceProfile without any specific executor request information, thus the executor information will utilize the default values from Default ResourceProfile (executor.cores=4).

The above code will require an extra executor which will have the same executor.cores/memory as the default ResourceProfile.

Different executor request information

from pyspark.resource import ExecutorResourceRequests, TaskResourceRequests, ResourceProfileBuilder

def filter_func(iterator):
    for pdf in iterator:
        yield pdf

df = spark.range(0, 100, 1, 4)

ereqs = ExecutorResourceRequests().cores(6)
treqs = TaskResourceRequests().cpus(5)
rp = ResourceProfileBuilder().require(treqs).require(ereqs).build
df.repartition(3).mapInArrow(lambda iter: iter, df.schema, False, rp).collect()

[wbo4958]

Hi @WeichenXu123 @zhengruifeng @tgravescs, Could you please help review this PR.

[wbo4958]

BTW, I will perform the similar manual tests for spark connect.

[tgravescs]

It would be nice to have more description here about what the challenge is (why it just didn't work) and then how you went about solving/fixing it.

Does this PR introduce any user-facing change?
No

Is this really true? doesn't this now allow user to call these together? Does it affect spark connect api since you mention that?

Please update description as it should be clear to reviewer without looking at the code what the proposed changes are.

[tgravescs]

def mapInArrow(
       self,
       func: "ArrowMapIterFunction",
       schema: Union[StructType, str],
       barrier: bool = False,
       profile: Optional[ResourceProfile] = None,

This is definitely a user facing api change

[wbo4958]

Hi @tgravescs, Thx for your comment.

I just updated the description and submitted a new commit which adds the unit tests to test the resource profile for sql dataframe and connect dataframe API.

@tgravescs @WeichenXu123 @zhengruifeng Could you help review this PR. Thx.

[wbo4958]

Hi @WeichenXu123, @tgravescs, @zhengruifeng, @Ngone51, I've separated the big PR into 2 PRs, one for sql, the other for connect. This PR is only for sql, Could you help review it.

[wbo4958]

@zhengruifeng Could you help review it again? Thx

[wbo4958]

Hi @HyukjinKwon, Could you help to review this PR. Thx

[WeichenXu123]

LGTM! :)

[HyukjinKwon]

I dont know much about this so I defer to @zhengruifeng and @WeichenXu123 who has more context.

[wbo4958]

Hi @HyukjinKwon, thx for your review and comment. The CI on the newest commit got passed, could you help review it again? Thx very much.

[HyukjinKwon]

I will defer to @WeichenXu123

[WeichenXu123]

LGTM

[wbo4958]

Hi @HyukjinKwon could you help merge?

[wbo4958]

Hi @WeichenXu123, @HyukjinKwon, @zhengruifeng, Could you help merge it? Thx very much.

[WeichenXu123]

Merged to master.