@SaraCorralLou
You can use Memory Profiler to profile your UDF. This will help us to understand which part of the UDF is leading to high memory utilization and has multiple calls during the execution.
https://www.databricks.com/blog/2022/11/30/memory-profiling-pyspark.html
