Parse PySpark string column of key-list dictionaries based on separate array column of keys

user_3773324 0
Viewed 587

I am currently stuck trying to extract the values in the 'keyValue' column below based on the ordered array of keys in the 'Keys' column below.

>>> df.select('reference_nbr', 'keyValue', 'Keys').show()
+--------------+-------------------------------------------------------------+---------------------+
|    ref_number|                                                     keyValue|                 Keys|
+--------------+-------------------------------------------------------------+---------------------+
|          AZQ5|{key39=[TBAX3, TBAX6, TBAXN], key46=[TBARO, TBAZ4, TBABN],...|[key1, key2, key3,...|
|          NXY3|{key39=[TBAX8, TBAX2, TBAXZ], key46=[TBARD, TBAZK, TBAX9],...|[key1, key2, key3,...|
|          QSW6|{key39=[TBAX5, TBAX3, TBAX8], key46=[TBARB, TBAZN, TBAX4],...|[key1, key2, key3,...|
|          LJB7|{key39=[TBAX3, TBAXN, TBAXL], key46=[TBARM, TBAZ2, TBAX3],...|[key1, key2, key3,...|
|          MKH9|{key39=[TBAX4, TBAX9, TBAXV], key46=[TBARB, TBAZB, TBAX1],...|[key1, key2, key3,...|
|          UFG1|{key39=[TBAX3, TBAX6, TBAXQ], key46=[TBARL, TBAZB, TBAX0],...|[key1, key2, key3,...|
|          WDE4|{key39=[TBAX6, TBAX7, TBAX9], key46=[TBARX, TBAX6, TBAX8],...|[key1, key2, key3,...|
|          VRX8|{key39=[TBAX3, TBAX1, TBAX0], key46=[TBARQ, TBAX9, TBAX3],...|[key1, key2, key3,...|
|          CIZ2|{key39=[TBAX3, TBAXC, TBAX2], key46=[TBARA, TBAXQ, TBAX1],...|[key1, key2, key3,...|
|          BEO3|{key39=[TBAX9, TBAXQ, TBAX4], key46=[TBARP, TBAXV, TBAX2],...|[key1, key2, key3,...|
+--------------+-------------------------------------------------------------+---------------------+
only showing top 20 rows

If I apply the UDF and withColumn() steps below, I can easily query the 'keyValue' column based on a specific key and insert the array of the key's values into a new column.

getKey4 = udf(lambda ar1: ar1.get('key4'))
df = df.withColumn("key4Values", getKey4(df["keyValue"]))

I am trying to perform the same steps as above but for every key in the order of the 'Keys' column. Desired output:

+--------------+-------------------------------------------------------------+---------------------+------------------------------------------------+
|    ref_number|                                                     keyValue|                 Keys|                                          Values|
+--------------+-------------------------------------------------------------+---------------------+------------------------------------------------+
|          AZQ5|{key39=[TBAX3, TBAX6, TBAXN], key46=[TBARO, TBAZ4, TBABN],...|[key1, key2, key3,...|[[TBAX4, TBAXQ, TBAXD],[TBAR1, TBAZA, TBABW],...|
|          NXY3|{key39=[TBAX8, TBAX2, TBAXZ], key46=[TBARD, TBAZK, TBAX9],...|[key1, key2, key3,...|[[TBAX5, TBAXA, TBAXC],[TBAR2, TBAZS, TBABE],...|
|          QSW6|{key39=[TBAX5, TBAX3, TBAX8], key46=[TBARB, TBAZN, TBAX4],...|[key1, key2, key3,...|[[TBAX6, TBAXZ, TBAXF],[TBAR3, TBAZD, TBABR],...|
|          LJB7|{key39=[TBAX3, TBAXN, TBAXL], key46=[TBARM, TBAZ2, TBAX3],...|[key1, key2, key3,...|[[TBAX7, TBAXC, TBAXG],[TBAR4, TBAZF, TBABT],...|
|          MKH9|{key39=[TBAX4, TBAX9, TBAXV], key46=[TBARB, TBAZB, TBAX1],...|[key1, key2, key3,...|[[TBAX8, TBAXV, TBAXH],[TBAR5, TBAZG, TBABY],...|
|          UFG1|{key39=[TBAX3, TBAX6, TBAXQ], key46=[TBARL, TBAZB, TBAX0],...|[key1, key2, key3,...|[[TBAX9, TBAXB, TBAXJ],[TBAR6, TBAZH, TBABU],...|
|          WDE4|{key39=[TBAX6, TBAX7, TBAX9], key46=[TBARX, TBAX6, TBAX8],...|[key1, key2, key3,...|[[TBAX0, TBAXN, TBAXK],[TBAR7, TBAZJ, TBABI],...|
|          VRX8|{key39=[TBAX3, TBAX1, TBAX0], key46=[TBARQ, TBAX9, TBAX3],...|[key1, key2, key3,...|[[TBAX2, TBAXM, TBAXL],[TBAR8, TBAZK, TBABO],...|
|          CIZ2|{key39=[TBAX3, TBAXC, TBAX2], key46=[TBARA, TBAXQ, TBAX1],...|[key1, key2, key3,...|[[TBAX3, TBAXA, TBAXO],[TBAR9, TBAZL, TBABP],...|
|          BEO3|{key39=[TBAX9, TBAXQ, TBAX4], key46=[TBARP, TBAXV, TBAX2],...|[key1, key2, key3,...|[[TBAX1, TBAXS, TBAXI],[TBAR0, TBAZQ, TBABZ],...|
+--------------+-------------------------------------------------------------+---------------------+------------------------------------------------+

I have tried the method below but I am getting the following error:

getKeyValues = udf(lambda ar1, ar2: {ar1.get(x) for x in ar2})
df.withColumn("Values", getKeyValues(df["keyValue"], df["Keys"])).show()

AttributeError: 'unicode' object has no attribute 'get'

I have tried a different UDF below that also gives the same unicode error:

def getKeyVals(ar1, ar2): 
    arr = []
    for x in ar2: 
        arr.append(ar1.get(x, None))
    return arr 
    
udf_split = udf(split, ArrayType(StringType()))

df.withColumn("test", udf_getKeyVals(df['keyValue'], df['Keys'])).show()

I have also tried the following function below but I am getting similar errors. This function comes from the following MungingData webpage

def working_fun(mapping):
    def f(ar1):
        for x in ar1:
            return mapping.get(x)
    return F.udf(f)

df.withColumn("test", working_fun(df["keyValue"])(F.col('Keys'))).show()

Would appreciate any tips or recommendations -- thanks!

Updated to include schema and Spark version details below

>>> df.select('reference_nbr', 'keyValue', 'Keys').schema.simpleString()
'struct<reference_nbr:string,keyValue:string,Keys:array<string>>'

>>> spark.version
u'2.3.2.3.1.5.6030-1'
pyspark apache-spark dataframe python
1 Answers

If we convert keyvalue column from string type to map type, we can use map_values function to extract values:

I used UDF to replace = in keyvalues by : so that we can change type to map using ast module.

from pyspark.sql import *
from pyspark.sql.functions import *
from pyspark.sql.types import *

spark = SparkSession.builder.master('local[*]').getOrCreate()


def get_dict(c):
    c = c.replace("=", ":")
    import ast
    dict_value = ast.literal_eval(c)
    return dict_value


get_dict_udf = udf(lambda c: get_dict(c), MapType(StringType(), ArrayType(IntegerType())))

# Sample dataframe
df = spark.createDataFrame(
    [('{"k3"= [6, 5, 4], "k1"= [4, 5, 1], "k8"= [8, 5, 6], "k5"= [7, 4, 3]}',
      ["k1", "k3", "k5", "k8"])]).toDF("keyvalue", "key")

df.withColumn("keyvalue", get_dict_udf("keyvalue")). \
    withColumn("values", sort_array(map_values("keyvalue"))).show(truncate=False)

+--------------------------------------------------------------------+----------------+--------------------------------------------+
|keyvalue                                                            |key             |values                                      |
+--------------------------------------------------------------------+----------------+--------------------------------------------+
|[k3 -> [6, 5, 4], k5 -> [7, 4, 3], k8 -> [8, 5, 6], k1 -> [4, 5, 1]]|[k1, k3, k5, k8]|[[4, 5, 1], [6, 5, 4], [7, 4, 3], [8, 5, 6]]|
+--------------------------------------------------------------------+----------------+--------------------------------------------+
user_87733090
Related