PySpark 使用 collect_list 收集不同长度的数组

Posted 2023-04-15

【中文标题】PySpark 使用 collect_list 收集不同长度的数组

【英文标题】：PySpark Using collect_list to collect Arrays of Varying Length

【发布时间】：2018-12-06 22:22:48

【问题描述】：

我正在尝试使用 collect_list 从两个不同的数据帧中收集数组（并保持顺序）。

Test_Data 和 Train_Data 具有相同的格式。

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

w = Window.partitionBy('Group').orderBy('date')

# Train_Data has 4 data points
# Test_Data has 7 data points
# desired target array:         [1, 1, 2, 3]
# desired MarchMadInd array:    [0, 0, 0, 1, 0, 0, 1]

sorted_list_diff_array_lens = train_data.withColumn('target', 
F.collect_list('target').over(w)
                                  )\
test_data.withColumn('MarchMadInd', F.collect_list('MarchMadInd').over(w))\
   .groupBy('Group')\
   .agg(F.max('target').alias('target'), 
    F.max('MarchMadInd').alias('MarchMadInd')
)

我意识到“test_data.withColumn”的语法不正确，但我想从 test_date 中选择 MarchMadInd 的数组，但 train_data 中的 em>target。所需的输出如下所示：

"target":[1, 1, 2, 3], "MarchMadInd":[0, 0, 0, 1, 0, 0, 1]

上下文：这是针对需要动态特征来包含预测周期的 DeepAR 时间序列模型（使用 AWS），但目标应该是历史数据。

【问题讨论】：

你需要一个键来连接两个数据框。

【参考方案1】：

解决方案涉及使用 pault 推荐的连接。

创建具有长度等于训练 + 预测周期的动态特征的数据框 创建一个数据帧，其目标值的长度正好等于训练周期。 使用 LEFT JOIN（动态特征数据位于 LEFT）将这些数据框组合在一起

现在，使用 collect_list 将创建所需的结果。