PySpark无需explode，如何筛选排序数组结构体列的子集？

我是PySpark新手，现有如下示例DataFrame：

df = spark.createDataFrame(
    (("7dc88", "D21", 14.14, 2, 10, [["msgA", 15, "a"],["msgB", 9, "g"],["msgC", 6, "z"],["msgD", 4, "m"],["msgE", 1, "e"]]),
     ("1c36a3", "D21", 32.14, 18, 45, [["msgA", 45, "n"],["msgB", 30, "q"],["msgC", 24, "h"],["msgD", 19, "y"],["msgE", 11, "c"]])),
     "uniqueId : string, tag : string, score : float, time0 : int, time1 : int, msgs : array<struct<msg : string, time : int, sysid : string>>")

DataFrame展示：

+--------+---+-----+-----+-----+---------------------------------------------------------------------------+
|uniqueId|tag|score|time0|time1|msgs                                                                       |
+--------+---+-----+-----+-----+---------------------------------------------------------------------------+
|7dc885  |D21|14.14|2    |10   |[{msgA, 15, a}, {msgB, 9, g}, {msgC, 6, z}, {msgD, 4, m}, {msgE, 1, e}]    |
|1c36a3  |D21|32.14|18   |45   |[{msgA, 45, n}, {msgB, 30, q}, {msgC, 24, h}, {msgD, 19, y}, {msgE, 11, c}]|
+--------+---+-----+-----+-----+---------------------------------------------------------------------------+

其中msgs列为array<struct>类型（包含msg、time、sysid字段）。实际场景中消息结构更复杂且含嵌套，每个uniqueId对应数百条消息，数据量达数百万条uniqueId。

time0和time1列定义了msgs子集的时间范围（包含边界），需要生成按time排序的新列msgs_subset，期望结果如下：

+--------+---+-----+-----+-----+------------------------------------------------------------+
|uniqueId|tag|score|time0|time1|msgs_subset                                                 |
+--------+---+-----+-----+-----+------------------------------------------------------------+
|7dc885  |D21|14.14|2    |10   |[{msgD, 4, m}, {msgC, 6, z}, {msgB, 9, g}]                  |
|1c36a3  |D21|32.14|18   |45   |[{msgD, 19, y}, {msgC, 24, h}, {msgB, 30, q}, {msgA, 45, n}]|
+--------+---+-----+-----+-----+------------------------------------------------------------+

我已通过以下代码实现该需求：

import pyspark.sql.functions as F

df_msg_subset = (
    df
    .withColumn("msg_explode", F.explode(F.col("msgs")))
    .withColumn("msg_time", F.col("msg_explode.time"))
    .filter((F.col("msg_time") >= F.col('time0')) & (F.col("msg_time") <= F.col('time1')))
    .sort(F.col("uniqueId"), F.col("msg_time"))
    .groupBy(list(set(df.columns) - {'msgs'})).agg(F.collect_list('msg_explode').alias('msgs_subset'))
)

问题 #

能否不使用explode实现该功能？我认为explode生成大量行再分组的方式性能开销较高，希望使用PySpark内置函数（而非UDF）直接操作msgs列实现，大致框架如下，但不知如何完善create_subset_column函数：

def create_subset_column(msgs, time0, time1):
    return F.sort_array(F.filter(...))

df_msg_subset = (
    df
    .withColumn("msgs_subset", create_subset_column("msgs", "time0", "time1"))
    .drop("msgs")
)

此外提供时间为毫秒级时间戳的示例DataFrame：

df = spark.createDataFrame(
    (
        ("7dc88", "D21", 14.14, 1642970972787, 1642970985027, [
            ["msgA", 1642970990067, "a"],
            ["msgB", 1642970985027, "g"],
            ["msgC", 1642970978077, "z"],
            ["msgD", 1642970972787, "m"],
            ["msgE", 1642970960897, "e"],
        ]),
        ("1c36a3", "D21", 32.14, 1642971056787, 1642971074107, [
            ["msgA", 1642971080687, "n"],
            ["msgB", 1642971074107, "q"],
            ["msgC", 1642971068777, "h"],
            ["msgD", 1642971062157, "y"],
            ["msgE", 1642971056787, "c"],
        ])
    ),
    "uniqueId:string, tag:string, score:float, time0:long, time1:long, msgs:array<struct<msg:string, time:long, sysid:string>>"
)

内容的提问来源于stack exchange，提问作者4mla1fn