Python dictionary lookup in PySpark

user_1062015 0
Viewed 244

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?

pyspark
1 Answers

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|
+------+------+------+--------------+
user_141657300
Related