Replace elements in an array with its corresponding number in pyspark

user_4452499 0
Viewed 2940

I have a dataframe shown below:

   +----------+--------------------------------+
   | Index    |           flagArray            |
   +----------+--------------------------------+
   |    1     | ['A','S','A','E','Z','S','S']  | 
   +----------+--------------------------------+
   |    2     | ['A','Z','Z','E','Z','S','S']  |
   +--------- +--------------------------------+

I want to represent array elements with its corresponding numeric values.

     A - 0
     F - 1
     S - 2
     E - 3
     Z - 4

So my output dataframe should look like 

   +----------+--------------------------------+--------------------------------+
   | Index    |           flagArray            |           finalArray           |
   +----------+--------------------------------+--------------------------------+
   |    1     | ['A','S','A','E','Z','S','S']  | [0, 2, 0, 3, 4, 2, 2]          | 
   +----------+--------------------------------+--------------------------------+
   |    2     | ['A','Z','Z','E','Z','S','S']  | [0, 4, 4, 3, 4, 2, 2]          |
   +--------- +--------------------------------+--------------------------------+

I have written an udf in pyspark where I am achieving it by writing some if else statements. Is there any better way to handle the same.

pyspark apache-spark-sql apache-spark python-3.x
4 Answers

For Spark 2.4+, you can simply use transform function to loop through each element of flagArray array and get its mapping value from a map column that you can create from that mapping using element_at:

mappings = {"A": 0, "F": 1, "S": 2, "E": 3, "Z": 4}
mapping_col = map_from_entries(array(*[struct(lit(k), lit(v)) for k, v in mappings.items()]))

df = df.withColumn("mappings", mapping_col) \
       .withColumn("finalArray", expr(""" transform(flagArray, x -> element_at(mappings, x))""")) \
       .drop("mappings")

df.show(truncate=False)
#+-----+---------------------+---------------------+
#|Index|flagArray            |finalArray           |
#+-----+---------------------+---------------------+
#|1    |[A, S, A, E, Z, S, S]|[0, 2, 0, 3, 4, 2, 2]|
#|2    |[A, Z, Z, E, Z, S, S]|[0, 4, 4, 3, 4, 2, 2]|
#+-----+---------------------+---------------------+
user_13865510

There doesn't seem to be a built-in function to map array elements, so here's perhaps an alternative udf, different from yours in that it uses a list comprehension:

dic = {'A':0,'F':1,'S':2,'E':3,'Z':4}
map_array = f.udf(lambda a: [dic[k] for k in a])
df.withColumn('finalArray', map_array(df['flagArray'])).show(truncate=False)

Output:

+------+---------------------+---------------------+
|Index |flagArray            |finalArray           |
+------+---------------------+---------------------+
|1     |[A, S, A, E, Z, S, S]|[0, 2, 0, 3, 4, 2, 2]|
|2     |[A, Z, Z, E, Z, S, S]|[0, 4, 4, 3, 4, 2, 2]|
+------+---------------------+---------------------+
user_57615580

For Spark 3.1+, you could call pyspark.sql.functions.transform and pyspark.sql.functions.element_at to do the job:

import pyspark.sql.functions as F

# create example DataFrame
df = spark.createDataFrame([
    (1, "A S A E Z S S".split(" ")),
    (2, "A Z Z E Z S S".split(" ")),
], ["Index", "flagArray"])

# make mapping column
mappings = {"A": 0, "F": 1, "S": 2, "E": 3, "Z": 4}
mapping_col = F.map_from_entries(F.array(*[F.struct(F.lit(k), F.lit(v)) for k, v in mappings.items()]))

# transform DataFrame
df = df.withColumn("mappings", mapping_col) \
       .withColumn("finalArray", F.transform("flagArray", lambda x: F.element_at("mappings", x))) \
       .drop("mappings")
user_61553630

After trying out multiple ways, I successfully created a VectorUDT column with DoubleType values.

import pyspark.sql.functions as F
import pyspark.sql.types as T
from pyspark.ml.functions import array_to_vector

def mapper_func(val):
r = []
for i, x in enumerate(val):
    v = str(mapping_dict.get(x, 0.0)) # Default replace with 0.0
    r.append(float(v))        
return r

my_udf = lambda z: mapper_func(z)
label_udf = F.udf(my_udf, T.ArrayType(T.DoubleType()))

df = df.withColumn('finalArray', array_to_vector(label_udf(df.flagArray)))

This will work even for values not present in the dictionary/map. Default replace is 0.0

user_68976030
Related