karthik_p
Databricks Partner

@Leo Bao​ Are you seeing this issue whenever you are getting different sizes of data sets, or your data set size is same. if issue you are seeing is due to larger dataset, please check below link and try to increase partition size Databricks Spark Pyspark RDD Repartition - "Remote RPC client disassociated. Likely due to container...

karthik.p