在火花数据框中的每一行的地图类型列中按键排序

Posted 2023-04-15

【中文标题】在火花数据框中的每一行的地图类型列中按键排序

【英文标题】：Sort by key in map type column for each row in spark dataframe

【发布时间】：2021-01-28 01:45:37

【问题描述】：

我有以下格式的 spark 数据框：

Name   LD_Value
A37    Map(10 -> 0.20,5 -> 0.30,17 -> 0.25)
A39    Map(11 -> 0.40,6 -> 0.67,24 -> 0.45)

我需要根据LD_Value 列中的键为每条记录按降序排序。

预期输出：

Name   LD_Value
A37    Map(17 -> 0.25,10 -> 0.20,5 -> 0.30)
A39    Map(24 -> 0.45,11 -> 0.40,6 -> 0.67)

是否可以对 spark 数据框中的地图类型列进行排序？

我研究了 spark 高阶函数，但没有运气。

【问题讨论】：

您可以定义一个 UDF 函数来对 maptype 列进行排序。看看这个链接：bmc.com/blogs/how-to-write-spark-udf-python

【参考方案1】：

你可以先使用map_keys函数获取map的keys，对keys数组进行排序，然后使用transform从原始map中获取每个key元素对应的值，最后通过创建更新map列使用map_from_arrays 函数从两个数组中创建一个新映射。

对于 Spark 3+，您可以使用比较器函数作为array_sort 函数的第二个参数，按降序对键数组进行排序：

from pyspark.sql import functions as F

df1 = df.withColumn(
    "LD_Value_keys",
    F.expr("array_sort(map_keys(LD_Value), (x, y) -> case when x > y then -1 when x < y then 1 else 0 end)")
).withColumn("LD_Value_values", F.expr("transform(LD_Value_keys, x -> LD_Value[x])")) \
 .withColumn("LD_Value", F.map_from_arrays(F.col("LD_Value_keys"), F.col("LD_Value_values"))) \
 .drop("LD_Value_keys", "LD_Value_values")

df1.show() 

#+----+----------------------------------+
#|Name|LD_Value                          |
#+----+----------------------------------+
#|A37 |[17 -> 0.25, 10 -> 0.2, 5 -> 0.3] |
#|A39 |[24 -> 0.45, 11 -> 0.4, 6 -> 0.67]|
#+----+----------------------------------+

对于 Spark ，您可以使用此 UDF 以降序对数组进行排序：

# array_sort_udf (array, reverse): if reverse = True then desc
array_sort_udf = F.udf(lambda arr, r: sorted(arr, reverse=r), ArrayType(StringType()))

并像这样使用它：

df.withColumn("LD_Value_keys", array_sort_udf(F.map_keys(F.col("LD_Value")), F.lit(True)))