[Python] PySpark with aggregations failing

I have a 106M dataframe with nested columns, that is, I have a few columns where the values are {[1,2,3,4,5], <timestamp>, 1, 2 ,3}. I'm trying to add a few more columns using when, then do an aggregation on the dataframe.

df_1 = df.withColumn("newCol", when((col("colA.field") == 1) & (col("colB.field1") == 2), col("colA.field1")).otherwise("colB.field1")))\
...
df_agg = df_1.groupby("colA","colB","colC","colD").agg(count(*).alias("numRecords"),
sort_array(collected_set("colE")).alias("colE"),
sum("colE").alias("colE"),
sum("colF").alias("colF"),
sum("colG").alias("colG"),
sum("colH").alias("colH"),
min("colI").alias("colI"),
max("colJ").alias("colJ"),
countDistinct("colK").alias("colK"),
first("colL").alias("colL"),
first("colM").alias("colM"),
first("colN").alias("colN"),
first("colO").alias("colO"),
sort_array(collected_set("colP")).alias("colP"),
sort_array(collected_set("colQ")).alias("colQ"),
max("colR").alias("colR"),
max("colS").alias("colS")
)


colL, colM, colN, colO are strings, and they are the same value for each group, so I simply want to get the first (or any) instance.

I've tried the following (separately) to simply do a df_agg.show(10, truncate=False), but have always gotten the error below.

Job aborted due to stage failure: ShuffleMapStage 4 (showString at NativeMethodAccessorImpl.java:0) has failed the maximum allowable number of times: 4. Most recent failure reason:
org.apache.spark.shuffle.FetchFailedException
at org.apache.spark.errors.SparkCoreErrors$.fetchFailedError(SparkCoreErrors.scala:132)
...
Caused by: org.apache.spark.ExecutorDeadException: The relative remote executor(Id: 253), which maintains the block data to fetch is dead.
at org.apache.spark.network.netty.NettyBlockTransferService$$anon$2.createAndStart(NettyBlockTransferService.scala:136)
...

1. Run the PySpark code on the original dataframe as it is with the nested columns (parquet files at 11.4GB)
2. Reduce the number of records from 106M to 99.8M using df.sample(0.943). I had previously successfully ran the same code on another similar dataframe with 99.9M rows (but without the nested columns, and the parquet files were at 5.8GB).
3. Flatten the schema and only selected relevant columns df_flat = df.select(col("colA.field1").alias("colA_field1"), ...)
4. Write the above dataframe df_flat into parquet files, start a new Spark session, read the parquet files back into df_flat before the step of adding more columns and aggregation. (parquet files at 4.4GB)

I also ran df.groupBy("colA", "colB", "colC", "colD").count().orderBy(desc("count")).show(), and the largest group has 68K records, followed by 37K, 27K, 21K, 13 groups with >10K records, and many more with ~9K or less. I think my data is skewed?

I'm not sure if the error I'm seeing is because of lack of Spark resources, or the way I'm doing the aggregation.

How can I optimize the code and get it to run successfully? This is only a small-scale test, and I would eventually need to run this on a much larger dataframe, in the order of billions of rows.

Continue reading...