mapping many keys to 1 values

Question

I am mapping countries in my data, I do the following for 1 key:value pair

from itertools import chain
countryDict = {
    'Australia': 'ANZ',
    'New Zealand': 'ANZ',
    'China': 'CHINA PLUS',
    'Hong Kong': 'CHINA PLUS', 
    'India': 'INDIA',
    'Japan': 'JAPAN',
    'Indonesia': 'SEA',
    'Malaysia': 'SEA',
    'Philippines': 'SEA',
    'Singapore': 'SEA',
    'Thailand': 'SEA',
    'VietNam': 'SEA',
    'South Korea': 'SOUTH KOREA',
    'Taiwan': 'TAIWAN',
    }

mapping_expr = F.create_map([F.lit(x) for x in chain(*countryDict.items())])
df = df.withColumn('Country_L1', mapping_expr[df['Customer_Country']])

I am not sure how to do it for multiple keys: 1 value, I am currently using when and otherwise to achieve but I want something like a dictionary to make it look simple

bbo = bbo.withColumn("Country_L1", F.when(F.col('customer_country').isin(
                'AU', 'CC', 'CK', 'CX', 'FJ', 'FM', 'GU', 'KI', 'MH', 'MP', 'NF', 'NR', 'NU', 'NZ', 'PG', 'PN', 'PW', 'TO', 'TV', 'WF', 'WS', 'SB', 'VU'), F.lit('ANZ')).otherwise(
                    F.when(F.col('customer_country').isin(
                        'CN', 'HK', 'MO'), F.lit('CHINA PLUS')).otherwise(
                            F.when(F.col('customer_country').isin(
                                'BD', 'BN', 'BT', 'ID', 'KH', 'KP', 'LA', 'LK', 'MM', 'MN', 'MV', 'MY', 'NP', 'PH', 'SG', 'TH', 'TL', 'VN'), F.lit('SEA')).otherwise(
                                    F.when(F.col('customer_country') == 'JP', F.lit('JAPAN')).otherwise(
                                        F.when(F.col('customer_country') == 'KR', F.lit('SOUTH KOREA')).otherwise(
                                            F.when(F.col('customer_country') == 'TW', F.lit('TAIWAN')).otherwise(
                                                F.lit('Non APAC'))
                                            )
                                        )
                                    ) 
                                )
                            )
                    )

any help to make this simple would be appreciated, thanks in advance.

Honestly, I have no clue what you are trying to achieve. Could you please provide a clear problem example with expected output and your (simpliefied) approach so far? — Jacob, Jul 29 '22 at 07:48
Hi @jacob, i am trying to group countries, in the 1st example it was minimal so it was easy to do via a dictionary but in latter example i am trying to simplify the iterative when conditions, i dont want to use it because i have many such mapping do and multiple when conditions become clumsy. if i could have a dictionary then it becomes easy to scale it, that is what i am asking help for. — Arun Mohan, Jul 29 '22 at 07:52
I still don't get it. What is the problem with creating a dictionary with `customer_country` as `key` and the corresponding category as `value`? Or is creating the dictionary in a _smarter_ way your question? — Jacob, Jul 29 '22 at 08:10
Or are you asking help to create the dictionary that looks like in your 1st example so that you can use your code in 1st example directly? — Jonathan Lam, Jul 29 '22 at 08:13

score 2 · Accepted Answer · answered Jul 29 '22 at 08:45

You can have a dict like this:

mapping = {
    'ANZ': ['AU', 'CC', 'CK', 'CX', 'FJ', 'FM', 'GU', 'KI', 'MH', 'MP', 'NF', 'NR', 'NU', 'NZ', 'PG', 'PN', 'PW', 'TO',
            'TV', 'WF', 'WS', 'SB', 'VU'],
    'CHINA PLUS': ['CN', 'HK', 'MO'],
    'SEA': ['BD', 'BN', 'BT', 'ID', 'KH', 'KP', 'LA', 'LK', 'MM', 'MN', 'MV', 'MY', 'NP', 'PH', 'SG', 'TH', 'TL', 'VN'],
    'JAPAN': ['JP'],
    'SOUTH KOREA': ['KR'],
    'TAIWAN': ['TW']
}

Then use python functools.reduce function to dynamically generate when expressions from the dict elements:

from functools import reduce
from pyspark.sql import functions as F

bbo = bbo.withColumn(
    "Country_L1",
    reduce(
        lambda a, b: a.when(F.col("customer_country").isin(b[1]), b[0]),
        mapping.items(),
        F
    ).otherwise('Non APAC')
)

Thank you very much for this, it works like a charm. can you plz explain why "F" in the reduce function? i am not able to understand. — Arun Mohan, Jul 29 '22 at 09:16
It is passed as the initial value for reduce function in order to do `F.when(...)` at the first iteration. Then, when expressions are simply chained (`when(...).when(...)...`). — blackbishop, Jul 29 '22 at 09:30

score 0 · Answer 2 · answered Jul 29 '22 at 08:16

look_up={ # Store the relations here
    'SEA': ['Indonesia','Malaysia','Philippines'],
    'JAPAN' :['Japan']
}
rev_lookup={value: key for (key, values) in look_up.items() for value in values } #Reverse key-value 
rev_lookup

Result:

{'Indonesia': 'SEA', 'Malaysia': 'SEA', 'Philippines': 'SEA', 'Japan': 'JAPAN'}

Possible duplicate with Reverse / invert a dictionary mapping

mapping many keys to 1 values

2 Answers2