[Python] Pyspark aggregations optimization

I have a huge dataframe with 3B rows. I'm running the PySpark code below with the Spark config.

spark = SparkSession\
.builder\
.appName("App")\
.config("spark.executor.memory","10g")\
.config("spark.executor.cores","4")\
.config("spark.executor.instances","6")\
.config("spark.sql.adaptive.enabled","true")\
.config("spark.dynamicAllocation.enabled","false")\
.enableHiveSupport()\
.getOrCreate()

df = spark.read.parquet("/data")
df = df.filter(col("colA").isNotNull() & col("colB").isNotNull())
df = df.withColumn("colK_udf",udf_function("colK"))

df_1 = df.withColumn("newCol", when((col("colA.field") == 1) & (col("colB.field1") == 2), col("colA.field1")).otherwise("colB.field1")))\
...
df_1 = df1.select(...)

df_agg = df_1.groupby("colA","colB","colC","colD").agg(count(*).alias("numRecords"),
sort_array(collected_set("colE")).alias("colE"),
sum("colF").alias("colF"),
sum("colG").alias("colG"),
sum("colH").alias("colH"),
sum("colL").alias("colL"),
min("colI").alias("colI"),
max("colJ").alias("colJ"),
countDistinct("colE").alias("colE"),
sort_array(collected_set("colP")).alias("colP"),
sort_array(collected_set("colQ")).alias("colQ"),
max("colR").alias("colR"),
max("colS").alias("colS")
)
df_agg.count()


I tried the code on a smaller dataframe with only 100M rows, and it worked. However, when I ran it on a dataframe with 3B rows, I get the error below when executing the last df_agg.count().

ERROR org.apache.spark.scheduler.TaskSchedulerImpl - Lost executor 1 on 2.2.2.2:
...
The API gave the following message: Pod ephemeral local storage usage exceeds the total limit of containers 50Gi.


I have already increased the pod local usage from 30GB to 50GB, but I can't increase it indefinitely. I'll keep getting this message for other executors, and the number of failed tasks just keeps increasing.

When I looked at the Spark UI, under the Details for the Job, I'll see Input going up to 50GB, and Shuffle Write going up to 150GB.

My data should be quite skewed too, with the top groups (after groupby) being much larger than the other groups.

I've tried to cache df_1 after select, but that didn't solve the problem.

What else can I try?

Continue reading...