Pyspark PandasUDF assignment

user_10296404 0
Viewed 68

I'm trying to evenly parallelise a workload across my workers via a PandasUDF.

The function is rather mem expensive, and I would like each worker to run the function at most n times simultaneously. Is there any way I can control how Pyspark distributes the PandasUDF across the workers to achieve this?

For example, let's say I have 20 groups in my Spark DataFrame, and my cluster consists of 5 workers, I would like each worker to execute the udf on 4 groups at the same time:

df.groupBy('GROUP_IX').applyInPandas(my_pandas_udf, some_output_schema)

I tried repartitioning by range the df beforehand, but that didn't give me the exact uniform workload distribution I'm aiming for (sometimes more than 4 partitions ended up in the same worker):

df.repartitionByRange(5, 'GROUP_IX').applyInPandas(my_pandas_udf, some_output_schema)

I understand this is not what Pyspark & PandasUDFs are intended for... but I'm trying to figure out if there's a way to do it.

Thanks!

pyspark apache-spark
1 Answers

Spark, by default, will execute as many UDFs per executor as many cores you assign to it, as per the executor.cores parameter.

If your application is memory bound you can try setting the executor.cores to 1.

Also, consider increasing the memory overhead as PandasUDF are executed outside the JVM spark memory space.

For an UDF application that is CPU bound I still don't have a way to allocate to each UDF more than one CPU.

user_118041050
Related