检查一列是不是与pyspark中的groupby连续

Posted 2023-04-15

【中文标题】检查一列是不是与pyspark中的groupby连续

【英文标题】：Check if a column is consecutive with groupby in pyspark检查一列是否与pyspark中的groupby连续

【发布时间】：2021-04-06 14:43:19

【问题描述】：

我有一个如下所示的 pyspark 数据框：

import pandas as pd
foo = pd.DataFrame('group': ['a','a','a','b','b','c','c','c'], 'value': [1,2,3,4,5,2,4,5])

我想创建一个新的二进制列is_consecutive，以指示value 列中的值是否与group 连续。

输出应如下所示：

foo = pd.DataFrame('group': ['a','a','a','b','b','c','c','c'], 'value': [1,2,3,4,5,2,4,5],
'is_consecutive': [1,1,1,1,1,0,0,0])

我如何在 pyspark 中做到这一点？

【参考方案1】：

您可以使用lag 将值与上一行进行比较并检查它们是否连续，然后使用min 确定给定组中的所有行是否连续。

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'consecutive', 
    F.coalesce(
        F.col('value') - F.lag('value').over(Window.partitionBy('group').orderBy('value')) == 1, 
        F.lit(True)
    ).cast('int')
).withColumn(
    'all_consecutive', 
    F.min('consecutive').over(Window.partitionBy('group'))
)

df2.show()
+-----+-----+-----------+---------------+
|group|value|consecutive|all_consecutive|
+-----+-----+-----------+---------------+
|    c|    2|          1|              0|
|    c|    4|          0|              0|
|    c|    5|          1|              0|
|    b|    4|          1|              1|
|    b|    5|          1|              1|
|    a|    1|          1|              1|
|    a|    2|          1|              1|
|    a|    3|          1|              1|
+-----+-----+-----------+---------------+

【参考方案2】：

您可以使用前导并减去与现有值相同的值，然后找到窗口的最大值，一旦完成，设置条件说 return 0 is max is >1 else return 1

w = Window.partitionBy("group").orderBy(F.monotonically_increasing_id())

(foo.withColumn("Diff",F.lead("value").over(w)-F.col("value"))
 .withColumn("is_consecutive",F.when(F.max("Diff").over(w)>1,0).otherwise(1))
.drop("Diff")).show()

---

+-----+-----+--------------+
|group|value|is_consecutive|
+-----+-----+--------------+
|    a|    1|             1|
|    a|    2|             1|
|    a|    3|             1|
|    b|    4|             1|
|    b|    5|             1|
|    c|    2|             0|
|    c|    4|             0|
|    c|    5|             0|
+-----+-----+--------------+