pandas udf type grouped map fails

user_b22ce5eeAl
New Contributor II

Hello,

I am trying to get the shap values for my whole dataset using pandas udf for each category of a categorical variable. It runs well when I run it on a few categories but when I want to run the function on the whole dataset my job fails. I see spills both on memory and disk and my shuffle read is around 40GB. I am not sure how to optimize my spark job here, I increased my cores to 160 and also Memory for both driver and workers but still not successful.

Any suggestion will be highly appreciated.

Thanks