karthik-kandiko
New Contributor II

I got to solve this by downgrading the Databricks runtime to 13.3 and had the below commands for optimization and it worked well in my case.

spark.conf.set("spark.sql.shuffle.partitions", "200")
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")