arrays - 用 PySpark 中的对应元素替换数组中的元素

Question

我有这个数据框:

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

我想用它们对应的数值来表示数组元素。

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

我在 PySpark 中编写了一个 udf，我通过编写一些 if else 语句来实现它。有没有更好的处理方法？

Answer 1

对于 Spark 2.4+，您可以简单地使用 transform函数循环遍历 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]|
#+-----+---------------------+---------------------+