在 Spark 中枚举连续相等值的块

Posted 2023-03-31

【中文标题】在 Spark 中枚举连续相等值的块

【英文标题】：Enumerate blocks of successively equal values in Spark

【发布时间】：2018-02-09 20:32:24

【问题描述】：

我想在 Spark DataFrame 中查找 true 组（或块）的 ID。也就是说，我想从这个出发：

>>> df.show()
+---------+-----+
|timestamp| bool|
+---------+-----+
|        1|false|
|        2| true|
|        3| true|
|        4|false|
|        5| true|
|        6| true|
|        7| true|
|        8| true|
|        9|false|
|       10|false|
|       11|false|
|       12|false|
|       13|false|
|       14| true|
|       15| true|
|       16| true|
+---------+-----+

到这里：

>>> df.show()
+---------+-----+-----+
|timestamp| bool|block|
+---------+-----+-----+
|        1|false|    0|
|        2| true|    1|
|        3| true|    1|
|        4|false|    0|
|        5| true|    2|
|        6| true|    2|
|        7| true|    2|
|        8| true|    2|
|        9|false|    0|
|       10|false|    0|
|       11|false|    0|
|       12|false|    0|
|       13|false|    0|
|       14| true|    3|
|       15| true|    3|
|       16| true|    3|
+---------+-----+-----+

（零是可选的，可以是 Null 或 -1 或任何更容易实现的）

【问题讨论】：

DataFrame 没有内在顺序，因此那里没有“块”。您需要一些订单定义，并使其在大型数据集上可行，需要一些分组变量。如果两者都有，请参阅下面链接的解决方案：

还有一列“timestamp”定义了订单

Pyspark : Custom window function的可能重复

另外，请查看How to make good reproducible Apache Spark Dataframe examples

@prcastro 我的解决方案对您有帮助吗？

【参考方案1】：

我在 scala 中有一个解决方案，应该很容易适应 pyspark。考虑以下数据框df：

+---------+-----+
|timestamp| bool|
+---------+-----+
|        1|false|
|        2| true|
|        3| true|
|        4|false|
|        5| true|
|        6| true|
|        7| true|
|        8| true|
|        9|false|
|       10|false|
|       11|false|
|       12|false|
|       13|false|
|       14| true|
|       15| true|
|       16| true|
+---------+-----+

那么你可以这样做：

df
  .withColumn("prev_bool",lag($"bool",1).over(Window.orderBy($"timestamp")))
  .withColumn("block",sum(when(!$"prev_bool" and $"bool",1).otherwise(0)).over(Window.orderBy($"timestamp")))
  .drop($"prev_bool")
  .withColumn("block",when($"bool",$"block").otherwise(0))
  .show()


+---------+-----+-----+
|timestamp| bool|block|
+---------+-----+-----+
|        1|false|    0|
|        2| true|    1|
|        3| true|    1|
|        4|false|    0|
|        5| true|    2|
|        6| true|    2|
|        7| true|    2|
|        8| true|    2|
|        9|false|    0|
|       10|false|    0|
|       11|false|    0|
|       12|false|    0|
|       13|false|    0|
|       14| true|    3|
|       15| true|    3|
|       16| true|    3|
+---------+-----+-----+