Fill NA in PySpark DataFrame by group with values from Pandas lookup table

I have a PySpark DataFrame with missing values in col2 that I would like to impute, based on values in col1. For example:

df
id    col1    col2
0      A       1
1      A      NA
2      B       2
3      B      NA
4      B       3

I would like to impute these missing values using a given Pandas lookup table:

pdf_lookup
id    col1    col2
0      A       4
1      B       5

So the desired result would be the following PySpark DataFrame:

id    col1    col2
0      A       1
1      A       4
2      B       2
3      B       5
4      B       3

What would be the most efficient way to do this? A scalable solution would be ideal since df may be very large with up to hundreds of columns (i.e. col3, ..., col500) that need to be imputed based on col1. Any suggestions would be appreciated!

You can do it with a join then a coalesce to keep only the non-null values of 2 columns:

pdf_lookup = pdf_lookup.select(col("col1"), col("col2").alias("col2_tmp"))
df.join(pdf_lookup, ["col1"], "left").withColumn("col2", coalesce(col("col2"), col("col2_tmp"))).drop("col2_tmp").show()