计数窗口中过滤值的出现

Question

我有以下数据框：

|     Timestamp     |   info   |
+-------------------+----------+
|2016-01-01 17:54:30|     8    |
|2016-02-01 12:16:18|     2    |
|2016-03-01 12:17:57|     1    |
|2016-04-01 10:05:21|     2    |
|2016-05-11 18:58:25|     7    |
|2016-06-11 11:18:29|     6    |
|2016-07-01 12:05:21|     3    |
|2016-08-11 11:58:25|     2    |
|2016-09-11 15:18:29|     9    |

我想创建一个名为 count 的新列，该列在窗口（-2，0）（当前行和前两个）中计数多少值是＆gt; 5（在我无法执行操作的前两个行中，我会输入0）。

最终的表应该是：

|     Timestamp     |   info   |   count  |
+-------------------+----------+----------+
|2016-01-01 17:54:30|     8    |     0    |
|2016-02-01 12:16:18|     2    |     0    |
|2016-03-01 12:17:57|     1    |     1    |
|2016-04-01 10:05:21|     2    |     0    |
|2016-05-11 18:58:25|     7    |     1    |
|2016-06-11 11:18:29|     6    |     2    |
|2016-07-01 12:05:21|     3    |     2    |
|2016-08-11 11:58:25|     2    |     1    |
|2016-09-11 15:18:29|     9    |     1    |

我尝试这样做，但它不起作用：

w = Window.orderBy('Timestamp').rowsBetween(-2, 0)
df_input = df_input.withColumn("count", F.when((F.count("info").over(w) > 5), F.count("info").over(w) > 5).otherwise(0))

Answer 1

如果您不介意针对前2行执行的计算，则以下内容将有效。

w = Window.orderBy('Timestamp').rowsBetween(-2, 0)
df_input = df_input.withColumn('count', F.count(F.when(F.col('info') > 5, 1)).over(w))

df_input.show()
# +-------------------+----+-----+
# |          Timestamp|info|count|
# +-------------------+----+-----+
# |2016-01-01 17:54:30|   8|    1|
# |2016-02-01 12:16:18|   2|    1|
# |2016-03-01 12:17:57|   1|    1|
# |2016-04-01 10:05:21|   2|    0|
# |2016-05-11 18:58:25|   7|    1|
# |2016-06-11 11:18:29|   6|    2|
# |2016-07-01 12:05:21|   3|    2|
# |2016-08-11 11:58:25|   2|    1|
# |2016-09-11 15:18:29|   9|    1|
# +-------------------+----+-----+

如果您需要2排为0，而无需更改窗口，则可以在条件时使用此：

w = Window.orderBy('Timestamp').rowsBetween(-2, 0)
df_input = df_input.withColumn(
    'count',
    F.when(F.size(F.collect_list('info').over(w)) == 3, F.count(F.when(F.col('info') > 5, 1)).over(w))
     .otherwise(0)
)
df_input.show()
# +-------------------+----+-----+
# |          Timestamp|info|count|
# +-------------------+----+-----+
# |2016-01-01 17:54:30|   8|    0|
# |2016-02-01 12:16:18|   2|    0|
# |2016-03-01 12:17:57|   1|    1|
# |2016-04-01 10:05:21|   2|    0|
# |2016-05-11 18:58:25|   7|    1|
# |2016-06-11 11:18:29|   6|    2|
# |2016-07-01 12:05:21|   3|    2|
# |2016-08-11 11:58:25|   2|    1|
# |2016-09-11 15:18:29|   9|    1|
# +-------------------+----+-----+