Have you tried enabling Apache Arrow in your job? This may improve memory utilization for your job. You can do that by adding this snippet to the top of your script or setting it as part of the Spark config for your job.
# Enable Arrow-based columnar data transfers
spark.conf.set("spark.sql.execution.arrow.enabled", "true")
See the docs here:
https://docs.databricks.com/spark/latest/spark-sql/spark-pandas.html#optimizing-conversion-between-s...
