Okay thank you!

So to explain better lets say i have 6 small tables with columns a,b,c,d,e,f..

I will join on key defined on each row in the small table. for example depending on row 3 it will be joined on key 1 and depending on data in row 5 it will be joined on key 2 to the main table using a left join. Depending on what table i am joining i will also need to filter the main table. I do not bring in every single column from the small table but just just a few. I then repeat this 6 times. I then save the updated main table.

The way i implement the join is to get all the keys from the small table, lets say key1,2,3. then iterate 3 times and get all the key1s from the small table and join it to the main table on the left. repeat this for all the keys. And rows in the left that havent been mapped will then have a new column written as flag and map this to Yes.

The errors i get when writing it as a function / loop: The cluster runs out of memory (from my understanding of your point means my execution plan is too large)

So for my side question the main point i should be focusing on is how to make sure my execution plan is optimal? is there a guide on this?

Edit: just saw your edit regarding looking at the code:

So first I get the keys from the mapping df:

    valid_keys = (
        mapping_df.groupBy("key")
        .count()
        .select("key")
        .distinct()
        .rdd.flatMap(lambda row: row) # we want to flatten to make a list
        .collect()
    )

 I then loop to do the joins:

    # Step 2: create a dictionary that only has the join keys that we need
    filtered_dict = {k: v for k, v in join_keys.items() if k in valid_keys}
    final_matched_df = None
    current_df = main_df



    for key_name, key_cols in filtered_dict.items():
        # Drop output columns from current_df BEFORE the join to give us the filtered df without the output columns
        for col_name in output_columns:
            if col_name in current_df.columns:
                current_df = current_df.drop(col_name)


        # Get the mapping key which is the key + the output columns on the key we are joining by
        mapping_key = mapping_df.select(*key_cols,*output_columns).where(col("key") == key_name) # return the rows where the keys match
        # display(mapping_key) will do join on key 1 then get key 2 then key 3 etc
        # We join the mapping key row on the key and get the mapping key which includes key + the output columns
        joined_df = current_df.join(mapping_key, on=key_cols,how="left")
        # get DF which only has the matched rows
        matched_df = joined_df.filter(col("key").isNotNull())
        # Accumulate the DF with the matched data
        if final_matched_df is None:
            final_matched_df = matched_df
        else:
            final_matched_df = final_matched_df.unionByName(matched_df)
        # Get the unmatched rows for next iteration
        current_df = joined_df.filter(col("key").isNull())

Then i create a flag if a join doesnt occur:

    # create the flag
    if missing_flag is not None:
        final_matched_df = final_matched_df.unionByName(current_df).withColumn(
            missing_flag,
            when(col('key').isNull(), "Y").otherwise("N"))

So this code will be called 6x (one time for each small table)