PySpark 如何对具有 2 个元素的结构数组进行排序

Question

我有以下架构

|-- node: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- score: string (nullable = true)
 |    |    |-- toID: string (nullable = true)

+-----------
 node|
+-------------
|[[57753, XV0912938...|
|[[51959, 02384848...|
|[[63487, 898898989...|

我想根据score和 select top 10 toID进行排序。 pyspark 的简单方法是什么？

Answer 1

sort_array(<array column>, asc=False) function 可用于对数组内的元素进行排序。 如果元素是结构，则数组根据结构中的第一个字段排序。 由于分数字段是结构中的第一个字段，您可以使用这个 function 直接对结构数组进行排序。

下面是一个根据结构中的分数获取前 5 个 ID 的示例。

假设您的数据如下所示

+----------------------------------------------------------------------------------------------------------------------------------------+
|arr_of_structs                                                                                                                          |
+----------------------------------------------------------------------------------------------------------------------------------------+
|[{270, foo0}, {237, foo1}, {2440, foo2}, {2059, foo3}, {3960, foo4}, {3009, foo5}, {842, foo6}, {1043, foo7}, {852, foo8}, {3925, foo9}]|
|[{1028, bar0}, {2571, bar1}, {1682, bar2}, {2543, bar3}, {689, bar4}, {11, bar5}, {1, bar6}, {740, bar7}, {113, bar8}, {701, bar9}]     |
+----------------------------------------------------------------------------------------------------------------------------------------+

root
 |-- arr_of_structs: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- _1: long (nullable = true)
 |    |    |-- _2: string (nullable = true)

代码会很简单

• 使用sort_array对结构数组进行排序
• 使用slice function 保留排序数组中的前N个元素。

data_sdf. \
    withColumn('arr_of_structs_top5', func.slice(func.sort_array('arr_of_structs', asc=False), 1, 5)). \
    withColumn('top5_id_only', func.expr('transform(arr_of_structs_top5, x -> x._2)')). \
    drop('arr_of_structs'). \
    show(truncate=False)

# +----------------------------------------------------------------------+------------------------------+
# |arr_of_structs_top5                                                   |top5_id_only                  |
# +----------------------------------------------------------------------+------------------------------+
# |[{3960, foo4}, {3925, foo9}, {3009, foo5}, {2440, foo2}, {2059, foo3}]|[foo4, foo9, foo5, foo2, foo3]|
# |[{2571, bar1}, {2543, bar3}, {1682, bar2}, {1028, bar0}, {740, bar7}] |[bar1, bar3, bar2, bar0, bar7]|
# +----------------------------------------------------------------------+------------------------------+