Performance issue with PySpark code

Hi amaaiia ,
Thanks for using Microsoft Fabric Community,

As mentioned by ObungiNiels, the performance challenges in your code arise from excessive iterative self-joins, which lead to high shuffle costs and memory exhaustion. The error message indicates that some executors are failing due to memory constraints or task failures.

The key factors affecting performance include the use of count(), which forces Spark to recompute in each iteration, overwriting df, which leads to unnecessary recalculations, and repeated self-joins, which cause excessive data shuffling, increasing execution time and memory usage.

Here's an optimized version of code:
Instead of overwriting df, you can accumulate results using unionByName(). Also, avoid count() by tracking new records using exceptAll().

from pyspark.sql.functions import col
# Sample Data
df = spark.createDataFrame([("AAA", 23656), ("BBB", 4647)], ["order", "value"])
df_a = spark.createDataFrame([
    ("AAA", "CCC", 6878),
    ("BBB", "DDD", 5676),
    ("DDD", "EEE", 5653),
    ("CCC", "FFF", 5666),
    ("EEE", "GGG", 54353),
    ("FFF", "HHH", 8878),
], ["order", "next_order", "next_value"])

# Initialize result DataFrame with the starting values
result_df = df
new_records = df
while new_records.count() > 0:  # Ensure loop only runs when new matches exist
    # Join only on new records to get next level data
    new_df = new_records.alias("r").join(
        df_a.alias("a"), col("r.order") == col("a.order"), "inner"
    ).select(
        col("a.next_order").alias("order"),
        col("a.next_value").alias("value")
    )

    # Remove already processed records to avoid re-processing
    new_records = new_df.exceptAll(result_df)

    # Accumulate new results without overwriting
    result_df = result_df.unionByName(new_records)

# Show final result
result_df.show()

This approach should help optimize performance and prevent memory issues. Let us know if you need further assistance.

If this post helps, then please consider Accept it as the solution to help the other members find it more quickly and a kudos would be appreciated.

Regards,
Vinay.