如何用PySpark按指定列表顺序将列值转为分组数组？

给定如下item_name列表：

[hpc_max,asset_min,off_median]

以及原始数据表：

需要按timestamp分组，将item_value按照指定的item_name顺序组成数组，得到如下结果：

如何用PySpark实现该需求？

推荐使用透视表+数组构造的方案，逻辑直观且性能稳定，步骤如下：

1. 初始化环境并定义目标顺序 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 初始化SparkSession
spark = SparkSession.builder.appName("OrderedItemValue").getOrCreate()

# 指定item_name的固定顺序
target_items = ["hpc_max", "asset_min", "off_median"]
# 构造结果中item_name列的固定数组表达式
item_order_expr = F.array(*[F.lit(item) for item in target_items])

2. 创建或加载原始DataFrame #

如果已有数据源，直接读取即可；这里用示例数据构造：

data = [
    ("hpc_max", 0.25, "2023-03-01T17:20:00.000+0000"),
    ("asset_min", 0.34, "2023-03-01T17:20:00.000+0000"),
    ("off_median", 0.30, "2023-03-01T17:30:00.000+0000"),
    ("hpc_max", 0.54, "2023-03-01T17:30:00.000+0000"),
    ("asset_min", 0.32, "2023-03-01T17:35:00.000+0000"),
    ("off_median", 0.67, "2023-03-01T17:20:00.000+0000"),
    ("asset_min", 0.54, "2023-03-01T17:30:00.000+0000"),
    ("off_median", 0.32, "2023-03-01T17:35:00.000+0000"),
    ("hpc_max", 0.67, "2023-03-01T17:35:00.000+0000")
]

df = spark.createDataFrame(data, schema=["item_name", "item_value", "timestamp"])

3. 按时间戳透视数据 #

通过pivot将每个item_name转为列，同时指定target_items保证列顺序与需求一致：

# 按timestamp分组，透视item_name列，聚合取每个item的对应value（每个时间戳下每个item唯一，用first/max均可）
pivoted_df = df.groupBy("timestamp").pivot("item_name", target_items).agg(F.first("item_value"))

4. 构造结果数组 #

将透视后的列按目标顺序组合成数组，同时添加固定的item_name列：

result_df = pivoted_df.select(
    item_order_expr.alias("item_name"),
    F.array(*[F.col(item) for item in target_items]).alias("item_value"),
    "timestamp"
).orderBy("timestamp")

5. 查看结果 #

result_df.show(truncate=False)

输出结果与需求完全匹配：

+--------------------------------+---------------+-----------------------------+
|item_name                       |item_value     |timestamp                    |
+--------------------------------+---------------+-----------------------------+
|[hpc_max, asset_min, off_median]|[0.25, 0.34, 0.67]|2023-03-01T17:20:00.000+0000|
|[hpc_max, asset_min, off_median]|[0.54, 0.54, 0.3]  |2023-03-01T17:30:00.000+0000|
|[hpc_max, asset_min, off_median]|[0.67, 0.32, 0.32]|2023-03-01T17:35:00.000+0000|
+--------------------------------+---------------+-----------------------------+

备选方案：自定义聚合函数 #

如果需要更灵活的逻辑（比如处理重复item的场景），可以用自定义聚合函数：

from pyspark.sql.types import StructType, StructField, StringType, DoubleType, ArrayType, MapType

class CollectOrderedValues(F.UserDefinedAggregateFunction):
    def inputSchema(self):
        return StructType([
            StructField("item_name", StringType()),
            StructField("item_value", DoubleType())
        ])
    
    def bufferSchema(self):
        return StructType([
            StructField("buffer", MapType(StringType(), DoubleType()))
        ])
    
    def dataType(self):
        return ArrayType(DoubleType())
    
    def deterministic(self):
        return True
    
    def initialize(self, buffer):
        buffer["buffer"] = {}
    
    def update(self, buffer, input):
        # 若同一时间戳下有重复item，保留最后一条的value
        buffer["buffer"][input[0]] = input[1]
    
    def merge(self, buffer1, buffer2):
        buffer1["buffer"].update(buffer2["buffer"])
    
    def evaluate(self, buffer):
        # 按目标顺序提取value
        return [buffer["buffer"].get(item, None) for item in target_items]

# 使用自定义聚合函数生成结果
result_df = df.groupBy("timestamp").agg(
    CollectOrderedValues().alias("item_value")
).select(
    item_order_expr.alias("item_name"),
    "item_value",
    "timestamp"
).orderBy("timestamp")

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