用 pyspark 中对应的数字替换数组中的元素答案

【问题标题】：Replace elements in an array with its corresponding number in pyspark用 pyspark 中对应的数字替换数组中的元素
【发布时间】：2020-06-15 13:04:20
【问题描述】：

我有一个如下所示的数据框：

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

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

所以我的输出数据框应该看起来像

   +----------+--------------------------------+--------------------------------+
   | 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 语句来实现它。有没有更好的处理方法。

【问题讨论】：

数组元素的顺序重要吗？
是的，订购很重要。否则会更简单
好吧，我也是这么想的。

标签： python-3.x apache-spark pyspark apache-spark-sql pyspark-dataframes

【解决方案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]|
#+-----+---------------------+---------------------+

【讨论】：

【解决方案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]|
+------+---------------------+---------------------+

【讨论】：

@ernes_k 我遇到了一个问题。当我尝试对 finalArray 列进行反向转换时，它的工作原理。而如果我的索引值都为 1 并且我正在执行 df.groupBy("Index").agg(max("finalArray").alias('flag')) 然后尝试对列标志进行反向转换，它不起作用
你想用agg(max("finalArray"))实现什么？
我的问题陈述是，对于所有具有公共索引的行，我必须派生一个数组，该数组将在行中的每个数组索引上具有最高优先级的标志。因此，我首先使用数字设置标志优先级，然后我使用索引进行 groupBy 并在每个数组索引上找到最大值。一旦我导出我的数组，我想恢复否。带有实际标志。

【解决方案3】：

对于 Spark 3.1+，您可以致电 pyspark.sql.functions.transform 和 pyspark.sql.functions.element_at 来完成这项工作：

import pyspark.sql.functions as F

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()]))

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

【讨论】：