pyspark, apply mapping

pyspark, apply mapping.py

# map values in pyspark
import pyspark.sql.functions as F
from itertools import chain

data = [['a', 1], ['b', 2], ['a', 3], ['d', 4]]
data = spark.createDataFrame(data, schema=['name', 'val'])
data.show()

# create mapping column
mapping = {'a': 'hello a', 'b': 'hello b', 'c': 'hello c'}
mapping = F.create_map([F.lit(x) for x in list(chain(*mapping.items()))])

# apply the mapping
res = data.select(mapping[F.col('name')])
res.show()
# |map(a, hello a, b, hello b, c, hello c)[name]|
# +---------------------------------------------+
# |                                      hello a|
# |                                      hello b|
# |                                      hello a|
# |                                         null|
# +---------------------------------------------+