Pyspark DataFrame taking too long to write

I have a list of ID to run through Market Basket Analysis model and the program will collect the respective 'rules' from their respective transactional data.

Program are as following step:

1. Get all transactional data
2. Get all unique ID
3. Run through the unique IDs in a loop

In the loop:

a. Filter transactional data by ID

b. Remove duplicates

c. Window function to combine items into "basket"

d. Declare model and fit data

e. Get Association rules dataframe

f. Append rules dataframe into a list

4. Declare a union function and merge all rules dataframes as a single dataframe
5. Write the results into a data storage

example:

df:

ID = ["001","002","003"]

rules:

Problem: The program took forever to write into data storage. Same goes for count and display. The steps before were okay, until I do those action mentioned.

I try to find out the root cause and determine this could be due to lazy evaluation of Spark. When calling agg, join, group by, it attaches these calls to the plan of the df object and creating a large execution plan. Then I convert my agg/groupby into a windows function and some other approach to solve it, but so far none has worked for me. Please help me out.

Here's the code

from pyspark.sql import functions

from pyspark.sql import functions as F
from pyspark.ml.fpm import FPGrowth
from datetime import datetime, timedelta
from delta.tables import *
from pyspark.sql.functions import *
from pyspark.sql.window import Window    

# Retrieve data    
df = somedata
rules_list = []
groups = [x[0] for x in df.select("ID").distinct().collect()]

windowSpec  = Window.partitionBy("BillNo").orderBy("ItemShortDescription")
windowSpecAgg  = Window.partitionBy("BillNo")

for subID in groups:
    print(subID)
    item_group_df = df.filter(col('SubscriptionID')==subID)
    basketdata = item_group_df.dropDuplicates(['BillNo', 'ItemShortDescription'])
    basketdata = basketdata.withColumn("row",row_number().over(windowSpec)).withColumn("Basket", F.collect_list("ItemShortDescription").over(windowSpecAgg)) \
    .where(col("row")==1).select("Basket")
    #Frequent Pattern Growth – FP Growth is a method of mining frequent itemsets using support, lift, and confidence.
    fpGrowth = FPGrowth(itemsCol="Basket", minSupport=0.003, minConfidence=0.005)
    model = fpGrowth.fit(basketdata)
    # generated association rules.
    rules = (model.associationRules).alias('rules')
    rules = rules.withColumn('antecedent', concat_ws(', ', 'antecedent')).withColumn('consequent', concat_ws(', ', 'consequent')).withColumn('SubscriptionId', lit(subID))
    rules_list.append(rules)

def unionAll(dfs):
    first, *_ = dfs  # Python 3.x, for 2.x you'll have to unpack manually
    return first.sql_ctx.createDataFrame(
        first.sql_ctx._sc.union([df.rdd for df in dfs]),
        first.schema
    )
rules_df = unionAll(rules_list)
write_path = f"/mnt/Rules"
rules_df.write.format('delta') \
       .partitionBy("ID") \
       .save(write_path)

Extra info:

rules_df.explain()

== Physical Plan ==
*(1) Scan ExistingRDD[antecedent#13576,consequent#13577,confidence#13578,lift#13579,support#13580,Id#13581]

Each rules DF from an ID is about 200-3000 rows. So in total its < 50k rows.

Image reference: High number of stages in writing task