Join vs mapping with a dictionary to add new data to a Pandas/PySpark dataframe?

user_15505624 0
Viewed 111

I have a big dataframe with a weather column that can take 5 different values (sunny, cloudy, rain, snow and other). I have to add another column with more information that depends solely on the weather value (eg. if it is a sunny day, add the value x to the new column, if it is a cloudy day add the value y, ...).

I'm wondering what is a "better" approach (more "standard" approach) for doing this. Either making a join with another small dataframe that has the weather column (join on this column) and a "new_data" column with these new values to be added (5 rows and 2 columns), or mapping each row of the dataframe with a dictionary with the weather values as keys and the new data to be added as values. An example of this last approach is shown below:

new_data = {"sunny": x, "cloudy": y, "rain": z, "snow": m, "other": n}

# Pandas
df["new_column"] = df["weather"].apply(lambda x: new_data[x])

# PySpark
from pyspark.sql.functions import udf, col
from pyspark.sql.types import FloatType
add_weather_data_f = udf(lambda x: new_data[x], FloatType())
df = df.withColumn("new_column", add_weather_data_f(col("weather")))

Which one is a "better" approach when working with a Pandas dataframe or a PySpark dataframe?

pyspark pandas python
1 Answers

Dont see a reason for a join here. Broadcast will be fast but I like to avoid joins wherever I can when working with big data.

from pyspark.sql import functions as F

new_data=(['sunny','x'],['cloud','y'],['rain','z'],['snow','m'],['other','n'])
new_data1=(F.when(F.col('weather')==(x[0]), F.lit(x[1])) for x in new_data)
df.withColumn('new_column', F.coalesce(*new_data1)).show(truncate=False)

+-------+----------+
|weather|new_column|
+-------+----------+
|sunny  |x         |
|cloud  |y         |
|rain   |z         |
|snow   |m         |
|other  |n         |
+-------+----------+
user_101073890
Related