python之将最大值时间戳放入 PySpark 的数组中

我有一个包含以下列的 PySpark 数据框(比如 df1)

1.> category - 包含独特的类别类型

2.> start_time_array - 按升序排列的时间戳数组

3.> end_time_array - 按升序排列的时间戳数组

4.> lenStart - start_time_array 中的数组长度

5.> lenEnd - end_time_array 中数组的长度

下面是df1的一个例子:

+--------+------------------------------------------+------------------------------------------+--------+------+ 
|category|                            end_time_array|                          start_time_array|lenStart|lenEnd| 
+--------+------------------------------------------+------------------------------------------+--------+------+ 
|       A|[2017-01-18 00:00:00, 2017-01-27 00:00:00]|[2017-01-16 00:00:00, 2017-01-25 00:00:00]|       2|     2| 
|       B|                     [2017-02-18 00:00:00]|[2017-02-14 00:00:00, 2017-02-21 00:00:00]|       2|     1| 
+--------+------------------------------------------+------------------------------------------+--------+------+ 

还有另一个数据框df2，它包含两列category 和timestamp。 df2 包含与 df1 相同的 category 值，并且 df1 中数组内时间戳的值是df2 中的时间戳。以下是 df2

+--------+-------------------+ 
|category|          timestamp| 
+--------+-------------------+ 
|       A|2017-01-16 00:00:00| 
|       A|2017-01-18 00:00:00| 
|       A|2017-01-25 00:00:00| 
|       A|2017-01-27 00:00:00| 
|       B|2017-02-14 00:00:00| 
|       B|2017-02-18 00:00:00| 
|       B|2017-02-21 00:00:00| 
|       B|2017-02-22 00:00:00| 
|       B|2017-02-24 00:00:00| 
|       B|2017-02-25 00:00:00| 
+--------+-------------------+ 

正如我们在上面的 df1 示例中看到的，对于 category -> B ，lenStart=2 不等于 lenEnd=1。在 df1 的所有行中，lenStart = lenEnd 或 lenStart = lenEnd+1 对于 df1 中的所有行 其中 lenStart = lenEnd+1，我想取 timestamp 的最大值(属于适当的 category)并将其放在end_time_array 中的数组。我该怎么做？

以下是使用 df2 中的信息处理 df1 后的预期输出

+--------+------------------------------------------+------------------------------------------+--------+------+ 
|category|                            end_time_array|                          start_time_array|lenStart|lenEnd| 
+--------+------------------------------------------+------------------------------------------+--------+------+ 
|       A|[2017-01-18 00:00:00, 2017-01-27 00:00:00]|[2017-01-16 00:00:00, 2017-01-25 00:00:00]|       2|     2| 
|       B|[2017-02-18 00:00:00, 2017-02-25 00:00:00]|[2017-02-14 00:00:00, 2017-02-21 00:00:00]|       2|     2| 
+--------+------------------------------------------+------------------------------------------+--------+------+ 

请您参考如下方法:

这应该适用于 Spark 1.5+:

import pyspark.sql.functions as F 
df3 = df1.where(F.col('lenStart') == (F.col('lenEnd') + 1)).select('category') 
df4 = df2.join(df3, 'Category').groupby('Category').agg(F.max('timestamp').alias('max')) 
df5 = df1.join(df4, 'Category', 'left') 
df1_changed = df5.withColumn('end_time_array', F.when(F.col('max').isNull(), 
    F.col('end_time_array')).otherwise(F.concat(F.col('end_time_array'), 
                                                F.array(F.col('max'))))) 
df1_changed = df1_changed.withColumn('lenEnd', F.size(F.col('end_time_array'))) 

df1_changed 将有一个修改过的 end_time_array 列，当您请求的条件适用时，它会添加所需的值，否则，它保持不变。