Заменить элементы в массиве на соответствующий номер в pyspark

Question

У меня есть датафрейм, показанный ниже:

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

Я хочу представить элементы массива с соответствующими им числовыми значениями c.

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

Поэтому мой выходной фрейм данных должен выглядеть следующим образом

   +----------+--------------------------------+--------------------------------+
   | 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]          |
   +--------- +--------------------------------+--------------------------------+

Я написал udf в pyspark, где я добился этого, написав некоторые операторы if else. Есть ли лучший способ справиться с тем же.

Answer 1

Для Spark 2.4+ вы можете просто использовать функцию transform, чтобы l oop через каждый элемент массива flagArray и получить его значение отображения из столбца карты, который вы можете создать из этого отображение с использованием 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]|
#+-----+---------------------+---------------------+

Answer 2

Кажется, что нет встроенной функции для отображения элементов массива, поэтому, возможно, здесь есть альтернативный udf, отличающийся от вашего тем, что он использует понимание списка:

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)

Вывод:

+------+---------------------+---------------------+
|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]|
+------+---------------------+---------------------+