[Python] PySpark program stuck at adding broadcast piece

I'm trying to write a PySpark program that filters for records in a very large dataframe (1-2B records) that matches some conditions on another smaller reference dataframe. This is done using a left join between the 2 dataframes, and then writing the results to a parquet file. When the reference dataframe is empty, the program runs successfully. But when the reference dataframe contains 414K records, the spark program hangs at the message storage.BlockManagerInfo: [dispatcher-BlockManagerMaster]: Added broadcast_4_piece0 in memory on xxx:12345 (size:4.0 MiB, free: 10 GiB).

My code is as follows.

def extract_to_df(spark, ref_db):
columns_to_drop = ["ColA", "ColB", "ColC"]

# Join conditions
join_cond_1 = (col("Col1") >= col("Col3a")) & (col("Col1") >= col("Col3b"))
join_cond_2 = (col("Col2") >= col("Col3a")) & (col("Col2") >= col("Col3b"))

df = spark.read.parquet(folder)
df_2 = df.filter(df["Col4"]=="abc").withColumn("Col1", udf_col(col("Col1a"))).withColumn("Col2", udf_col(col("Col2a")))

df_tmp = df_2.join(ref_db, on=join_cond_1, how="left").drop(*columns_to_drop).withColumnRenamed("Col5", "Col5a")
df_results = df_tmp.join(ref_db, on=join_cond_2, how="left").drop(*columns_to_drop).withColumnRenamed("Col6", "Col6a")
df_final_results = df_results.dropna(subset=["Col5a", "Col6a"])

df_final_results.write.mode("overwrite").parquet(output_folder)

def main():
ref_db = spark.read.parquet("/ref_db.parquet")
ref_db.persist()

extract_to_df(spark, ref_db)

if __name__ == "__main__":
main()


What is wrong with the code?

Continue reading...