Pyspark / replacement based on a different column string

Viewed 33
location = ['Seoul', 'Busan', 'Jeju']

df = df.withColumn('new_key_location', 
                    F.when(F.col('key_location').isin(location),
                    F.col('key_location')).otherwise(F.lit('Others')))

I created a new column called 'new_key_location', and except for locations Seoul, Busan, and Jeju, I named other locations as others.

However, I have a different column called 'location name'. There are also Seoul, Busan, and Jeju, but they were named as others in the new_key_location column.

How can I change 'others' in the new_key_location(newly created column) based on Seoul, Busan, and Jeju in the column 'location name.'

key Location location name new_key_location
Others Seoul Others → Seoul
Others Others Others
Seoul Seoul Seoul
Others Busan Others → Busan
1 Answers

A simple case when (when().otherwise()) should work.

data_sdf. \
    withColumn('new_location',
               func.when((func.col('new_key_location') == 'Others') & 
                         (func.col('location_name').isin(location)), func.col('location_name')).
               otherwise(func.col('new_key_location'))
               )

If new_key_location is "Others" but location_name is one of the locations in the list, it will update the value from "Others" to the one in location_name column.

Related