'if in inside of UDF'

user_13520226 0
Viewed 40

I have a table like this:

+----------------------+
|Country_state         |
+----------------------+
| Virginia             |
| New Jersey           |
| British Columbia     |
|Over the North Sea    |
| Germany              |
| Belgium              |
| Germany              |
| Bulgeria             |
| England              |
| England              |
| Germany              |
| England              |
| Belgium              |
...

and i need to get countries so i wrote simple udf:

def USA(co):
    states = ["Alaska", "Alabama", "Arkansas", "American Samoa", "Arizona", "California", "Colorado", "Connecticut", "District ", "of Columbia", "Delaware", "Florida", "Georgia", "Guam", "Hawaii", "Iowa", "Idaho", "Illinois", "Indiana", "Kansas", "Kentucky", "Louisiana", "Massachusetts", "Maryland", "Maine", "Michigan", "Minnesota", "Missouri", "Mississippi", "Montana", "North Carolina", "North Dakota", "Nebraska", "New Hampshire", "New Jersey", "New Mexico", "Nevada", "New York", "Ohio", "Oklahoma", "Oregon", "Pennsylvania", "Puerto Rico", "Rhode Island", "South Carolina", "South Dakota", "Tennessee", "Texas", "Utah", "Virginia", "Virgin Islands", "Vermont", "Washington", "Wisconsin", "West Virginia", "Wyoming"]

    if co in states:
        return "USA"
    else:
        return co

but this if seems to always be false and i don't know why.

This is how i'm calling it:

usa = udf(USA, StringType())
finalCountry = c. withColumn("CountryFINAL", usa(c.Country_state))
pyspark python
1 Answers

No need for UDF, use .isin and when/otherwise

from pyspark.sql import functions as F

states = ["Alaska", "Alabama", "Arkansas", "American Samoa", "Arizona", "California", "Colorado", "Connecticut", "District ", "of Columbia", "Delaware", "Florida", "Georgia", "Guam", "Hawaii", "Iowa", "Idaho", "Illinois", "Indiana", "Kansas", "Kentucky", "Louisiana", "Massachusetts", "Maryland", "Maine", "Michigan", "Minnesota", "Missouri", "Mississippi", "Montana", "North Carolina", "North Dakota", "Nebraska", "New Hampshire", "New Jersey", "New Mexico", "Nevada", "New York", "Ohio", "Oklahoma", "Oregon", "Pennsylvania", "Puerto Rico", "Rhode Island", "South Carolina", "South Dakota", "Tennessee", "Texas", "Utah", "Virginia", "Virgin Islands", "Vermont", "Washington", "Wisconsin", "West Virginia", "Wyoming"]

df.withColumn("countryFINAL", F.when(F.col("Country_state").isin(states), F.lit("USA"))\
                               .otherwise(F.col("Country_state"))).show()

#+------------------+------------------+
#|     Country_state|      countryFINAL|
#+------------------+------------------+
#|          Virginia|               USA|
#|        New Jersey|               USA|
#|  British Columbia|  British Columbia|
#|Over the North Sea|Over the North Sea|
#|           Germany|           Germany|
#|           Belgium|           Belgium|
#|           Germany|           Germany|
#|          Bulgeria|          Bulgeria|
#|           England|           England|
#|           England|           England|
#|           Germany|           Germany|
#|           England|           England|
#|           Belgium|           Belgium|
#+------------------+------------------+
user_101073890
Related