Pyspark-使用两个时间索引进行窗口功能

Question

我有一个数据框，其中每一行都有两个日期列。 我想创建一个窗口函数，其范围介于特定日期范围内的行数，两个日期列均在该范围内。 在下面的情况下，行的两个时间戳都必须在当前行的时间戳之前，以包括在计数中。

示例df包括count列：

    +---+-----------+-----------+-----+
    | ID|Timestamp_1|Timestamp_2|Count|
    +---+-----------+-----------+-----+
    |  a|          0|          3|    0|
    |  b|          2|          5|    0|
    |  d|          5|          5|    3|
    |  c|          5|          9|    3|
    |  e|          8|         10|    4|
    +---+-----------+-----------+-----+

我尝试创建两个窗口并在这两个窗口上创建新列：

    w_1 = Window.partitionBy().orderBy('Timestamp_1').rangeBetween(Window.unboundedPreceding, 0)
    w_2 = Window.partitionBy().orderBy('Timestamp_2').rangeBetween(Window.unboundedPreceding, 0)

    df = df.withColumn('count', F.count('ID').over(w_1).over(w_2))

但是，Pyspark不允许这样做，因此会出现错误。

有任何想法吗？ SQL解决方案也很好！

Answer 1

自我参与会成功吗？

from pyspark.sql import functions as F

df_count = (
    df.alias('a')
    .join(
        df.alias('b'),
        (F.col('b.Timestamp_1') <= F.col('a.Timestamp_1')) &
        (F.col('b.Timestamp_2') <= F.col('a.Timestamp_2')),
        'left'
    )
    .groupBy(
        'a.ID'
    )
    .agg(
        F.count('b.ID').alias('count')
    )
)

df = df.join(df_count, 'ID')