Python dictionary lookup in PySpark - pyspark

Question

Struggling with the following in PySpark. I have a dictionary in Python that looks like this:
COUNTRY_MAP = {
"AND": "AD", "ARE": "AE", "AFG": "AF", "ATG": "AG", "AIA": "AI", ... };
I now want to build up a value consisting of 3 columns, say value1, value2 and value3. The problem is that value3 needs to use the above lookup to convert the 3-letter code to a 2-letter code and if it does not exist, then "NONE" should be used, i.e.
from pyspark.sql import functions as sf
combined = sf.trim(sf.concat(sf.col("value1"), sf.lit(":"), sf.col("value2"), sf.lit(":"),
sf.coalesce(sf.col("value3"), "NONE")))
tmp = (df.withColumn('COMBINED_FIELD', combined)
...<other stuff>
)
This gives me values like "abc:4545:AND", "def:7789:ARE" and "ghi:1122:NONE". I now need: "abc:4545:AD", "def:7789:AE" and "ghi:1122:NONE".
As a newbie in PySpark, I am really struggling to get this working. Do you know?

Answer 1

You can convert the dictionary into a map type column and get the values using value3 as the key:
import pyspark.sql.functions as F
COUNTRY_MAP = {"AND": "AD", "ARE": "AE", "AFG": "AF", "ATG": "AG", "AIA": "AI"}
result = df.withColumn(
'combined_field',
F.trim(
F.concat_ws(':',
'value1', 'value2',
F.coalesce(
F.create_map(*sum([[F.lit(k), F.lit(v)] for (k,v) in COUNTRY_MAP.items()], []))[F.col('value3')],
F.lit('NONE')
)
)
)
)
result.show()
+------+------+------+--------------+
|value1|value2|value3|combined_field|
+------+------+------+--------------+
| abc| 4545| AND| abc:4545:AD|
| def| 7789| ARE| def:7789:AE|
| ghi| 1122| NONE| ghi:1122:NONE|
+------+------+------+--------------+