OK, without having your code or DAG, it's a little difficult to figure this out.  But here's something that should work.  First, figure out who many Memory Partitions you have.  Apparently, your Memory Partitions are too big for the cluster, hence the OOM. Use this generic code as a template. 

num_partitions = df.rdd.getNumPartitions() print(num_partitions)