Hi @fh, ensure there are no redundant calls in your code and use logging within the function to track its execution. Check your cluster configuration and data partitioning to avoid misconfigurations and uneven distribution, respectively. Verify that your Databricks Runtime (DBR 15.2 ML) isn't causing the problem, and experiment with smaller data batches. 

I hope this helps! Let me know if you have any other questions or need further assistance.