map values in a dataframe from a dictionary using pyspark

Question

I want to know how to map values in a specific column in a dataframe.

I have a dataframe which looks like:

df = sc.parallelize([('india','japan'),('usa','uruguay')]).toDF(['col1','col2'])

+-----+-------+
| col1|   col2|
+-----+-------+
|india|  japan|
|  usa|uruguay|
+-----+-------+

I have a dictionary from where I want to map the values.

dicts = sc.parallelize([('india','ind'), ('usa','us'),('japan','jpn'),('uruguay','urg')])

The output I want is:

+-----+-------+--------+--------+
| col1|   col2|col1_map|col2_map|
+-----+-------+--------+--------+
|india|  japan|     ind|     jpn|
|  usa|uruguay|      us|     urg|
+-----+-------+--------+--------+

I have tried using the lookup function but it doesn't work. It throws error SPARK-5063. Following is my approach which failed:

def map_val(x):
    return dicts.lookup(x)[0]

myfun = udf(lambda x: map_val(x), StringType())

df = df.withColumn('col1_map', myfun('col1')) # doesn't work
df = df.withColumn('col2_map', myfun('col2')) # doesn't work

Answer 1

I think the easier way is just to use a simple dictionary and df.withColumn .

from itertools import chain
from pyspark.sql.functions import create_map, lit

simple_dict = {'india':'ind', 'usa':'us', 'japan':'jpn', 'uruguay':'urg'}

mapping_expr = create_map([lit(x) for x in chain(*simple_dict.items())])

df = df.withColumn('col1_map', mapping_expr[df['col1']])\
       .withColumn('col2_map', mapping_expr[df['col2']])

df.show(truncate=False)

Answer 2

udf way

I would suggest you to change the list of tuples to dicts and broadcast it to be used in udf

dicts = sc.broadcast(dict([('india','ind'), ('usa','us'),('japan','jpn'),('uruguay','urg')]))

from pyspark.sql import functions as f
from pyspark.sql import types as t
def newCols(x):
    return dicts.value[x]

callnewColsUdf = f.udf(newCols, t.StringType())

df.withColumn('col1_map', callnewColsUdf(f.col('col1')))\
    .withColumn('col2_map', callnewColsUdf(f.col('col2')))\
    .show(truncate=False)

which should give you

+-----+-------+--------+--------+
|col1 |col2   |col1_map|col2_map|
+-----+-------+--------+--------+
|india|japan  |ind     |jpn     |
|usa  |uruguay|us      |urg     |
+-----+-------+--------+--------+

join way (slower than udf way)

All you have to do is change the dicts rdd to dataframe too and use two joins with aliasings as following

df = sc.parallelize([('india','japan'),('usa','uruguay')]).toDF(['col1','col2'])

dicts = sc.parallelize([('india','ind'), ('usa','us'),('japan','jpn'),('uruguay','urg')]).toDF(['key', 'value'])

from pyspark.sql import functions as f
df.join(dicts, df['col1'] == dicts['key'], 'inner')\
    .select(f.col('col1'), f.col('col2'), f.col('value').alias('col1_map'))\
    .join(dicts, df['col2'] == dicts['key'], 'inner') \
    .select(f.col('col1'), f.col('col2'), f.col('col1_map'), f.col('value').alias('col2_map'))\
    .show(truncate=False)

which should give you the same result

Answer 3

Similar to Ali AzG, but pulling it all out into a handy little method if anyone finds it useful

from itertools import chain
from pyspark.sql import DataFrame
from pyspark.sql import functions as F
from typing import Dict

def map_column_values(df:DataFrame, map_dict:Dict, column:str, new_column:str="")->DataFrame:
    """Handy method for mapping column values from one value to another

    Args:
        df (DataFrame): Dataframe to operate on 
        map_dict (Dict): Dictionary containing the values to map from and to
        column (str): The column containing the values to be mapped
        new_column (str, optional): The name of the column to store the mapped values in. 
                                    If not specified the values will be stored in the original column

    Returns:
        DataFrame
    """
    spark_map = F.create_map([F.lit(x) for x in chain(*map_dict.items())])
    return df.withColumn(new_column or column, spark_map[df[column]])

This can be used as follows

from pyspark.sql import Row, SparkSession
spark = SparkSession.builder.master("local[3]").getOrCreate()
df = spark.createDataFrame([Row(A=0), Row(A=1)])
df = map_column_values(df, map_dict={0:"foo", 1:"bar"}, column="A", new_column="B")
df.show()
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
#+---+---+
#|  A|  B|
#+---+---+
#|  0|foo|
#|  1|bar|
#+---+---+