Question: PySpark: Create a new column in dataframe based on another dataframe's cell values

I have PySpark dataframe dhl_price of the following form:
+------+-----+-----+-----+------+
|Weight|    A|    B|    C|     D|
+------+-----+-----+-----+------+
|     1|16.78|17.05|20.23|  40.1|
|     2|16.78|17.05|20.23| 58.07|
|     3|18.43|18.86| 25.0| 66.03|
|     4|20.08|20.67|29.77| 73.99|

So you can get the delivery price based on the category (i.e. the columns A, B, C, D) and the weight of your parcel (i.e. the first column Weight) and for weights larger than 30, we have prices specified only for 30, 40, 50 etc.
I also have PySpark dataframe requests, one row for each request by a customer. It includes columns product_weight, Type (the category that is in dhl_price). I want to create a new column in requests delivery_fee based on dhl_price dataframe. In particular, for each row in dhl_price column, I want to get a cell value in dhl_price where column is the one specified in column Type and row is the one specified in column product_weight of requests dataframe.
So far I could code it in pandas:
def get_dhl_fee(weight, type):
    if weight <= 30:
        price = dhl_price.loc[dhl_price["Weight"] == weight][type].values[0]
    else:
        price = dhl_price.loc[dhl_price["Weight"] >= weight].reset_index(drop = True).iloc[0][type].values[0]
    return price
    
new_requests["dhl_fee"] = new_requests.apply(lambda x: get_dhl_fee(x["product_weight_g"], x["Type"]), axis = 1)

How can I do the same with PySpark? I tried to use PySpark's UDF:
# Define the UDF (User-Defined Function) for calculating DHL fee
@fn.udf(returnType=DoubleType())
def get_dhl_fee(product_weight_g, calculate_way):
    broadcast_dhl_price = fn.broadcast(dhl_price)

    if weight <= 30:
        price = broadcast_dhl_price.filter(dhl_price["Weight"] == weight).select(calculate_way).first()[0]
    else:
        price = broadcast_dhl_price.filter(dhl_price["Weight"] >= weight).select(calculate_way).first()[0]

    return price

# Register the UDF
sc.udf.register("get_dhl_fee", get_dhl_fee)

# Apply the UDF to calculate dhl_fee column
requests = requests.withColumn("dhl_fee", get_dhl_fee(fn.col("product_weight"), fn.col("Type")))

but it returns error SPARK-5063:

"It appears that you are attempting to reference SparkContext from a
broadcast variable, action, or transformation. SparkContext can only
be used on the driver, not in code that it run on workers."



Answer:
Setup
dhl_price.show()
+------+-----+-----+-----+-----+
|Weight|    A|    B|    C|    D|
+------+-----+-----+-----+-----+
|     1|16.78|17.05|20.23| 40.1|
|     2|16.78|17.05|20.23|58.07|
|     3|18.43|18.86| 25.0|66.03|
|     4|20.08|20.67|29.77|73.99|
|    30|20.08|20.67|29.77|73.99|
|    40|21.08|21.67|30.77|74.99|
+------+-----+-----+-----+-----+

requests.show()
+--------------+----+
|product_weight|type|
+--------------+----+
|             1|   B|
|             2|   D|
|             4|   A|
|            30|   A|
|           100|   C|
+--------------+----+

Code
Create a column of map type which maps type to corresponding price for a given weight
c = dhl_price.columns[1:]
dhl_price_map = dhl_price.select('Weight', F.map_from_arrays(F.array(*map(F.lit, c)), F.array(*c)).alias('price'))
dhl_price_map.show()


+------+--------------------+
|Weight|               price|
+------+--------------------+
|     1|{A -> 16.78, B ->...|
|     2|{A -> 16.78, B ->...|
|     3|{A -> 18.43, B ->...|
|     4|{A -> 20.08, B ->...|
|    30|{A -> 20.08, B ->...|
|    40|{A -> 21.08, B ->...|
+------+--------------------+

Assign a unique identifier for each row in the requests dataframe
df = requests.withColumn('id_', F.monotonically_increasing_id())
df.show()
+--------------+----+-----------+
|product_weight|type|        id_|
+--------------+----+-----------+
|             1|   B| 8589934592|
|             2|   D|25769803776|
|             4|   A|34359738368|
|            30|   A|51539607552|
|           100|   C|60129542144|
+--------------+----+-----------+

Join the two dataframes on weight condition then use the indexing to yank the value of price corresponding to type for each row
df = df.join(dhl_price_map, on=df['product_weight'] >= dhl_price_map['Weight'], how='left')
df = df.withColumn('dhl_fee', F.expr("price[type]"))
df.show()


+--------------+----+-----------+------+--------------------+-------+
|product_weight|type|        id_|Weight|               price|dhl_fee|
+--------------+----+-----------+------+--------------------+-------+
|             1|   B| 8589934592|     1|{A -> 16.78, B ->...|  17.05|
|             2|   D|25769803776|     1|{A -> 16.78, B ->...|   40.1|
|             2|   D|25769803776|     2|{A -> 16.78, B ->...|  58.07|
|             4|   A|34359738368|     1|{A -> 16.78, B ->...|  16.78|
|             4|   A|34359738368|     2|{A -> 16.78, B ->...|  16.78|
|             4|   A|34359738368|     3|{A -> 18.43, B ->...|  18.43|
|             4|   A|34359738368|     4|{A -> 20.08, B ->...|  20.08|
|            30|   A|51539607552|     1|{A -> 16.78, B ->...|  16.78|
|            30|   A|51539607552|     2|{A -> 16.78, B ->...|  16.78|
|            30|   A|51539607552|     3|{A -> 18.43, B ->...|  18.43|
|            30|   A|51539607552|     4|{A -> 20.08, B ->...|  20.08|
|            30|   A|51539607552|    30|{A -> 20.08, B ->...|  20.08|
|           100|   C|60129542144|     1|{A -> 16.78, B ->...|  20.23|
|           100|   C|60129542144|     2|{A -> 16.78, B ->...|  20.23|
|           100|   C|60129542144|     3|{A -> 18.43, B ->...|   25.0|
|           100|   C|60129542144|     4|{A -> 20.08, B ->...|  29.77|
|           100|   C|60129542144|    30|{A -> 20.08, B ->...|  29.77|
|           100|   C|60129542144|    40|{A -> 21.08, B ->...|  30.77|
+--------------+----+-----------+------+--------------------+-------+

Create a window specification to rank the Weight per unique row in the requests, and then filter the rows to retain only those corresponding to the maximum weight in dhl_price.
W = Window.partitionBy('id_').orderBy(F.desc('Weight'))
df = df.withColumn('rank', F.dense_rank().over(W)).filter('rank == 1')

+--------------+----+-----------+------+--------------------+-------+----+
|product_weight|type|        id_|Weight|               price|dhl_fee|rank|
+--------------+----+-----------+------+--------------------+-------+----+
|             1|   B| 8589934592|     1|{A -> 16.78, B ->...|  17.05|   1|
|             2|   D|25769803776|     2|{A -> 16.78, B ->...|  58.07|   1|
|             4|   A|34359738368|     4|{A -> 20.08, B ->...|  20.08|   1|
|            30|   A|51539607552|    30|{A -> 20.08, B ->...|  20.08|   1|
|           100|   C|60129542144|    40|{A -> 21.08, B ->...|  30.77|   1|
+--------------+----+-----------+------+--------------------+-------+----+

Drop the extra columns
df = df.select(*requests.columns, 'dhl_fee')
df.show()

+--------------+----+-------+
|product_weight|type|dhl_fee|
+--------------+----+-------+
|             1|   B|  17.05|
|             2|   D|  58.07|
|             4|   A|  20.08|
|            30|   A|  20.08|
|           100|   C|  30.77|
+--------------+----+-------+