如何对PySpark中的分组数据应用条件计数（带重置）？

我有一个PySpark代码，它可以有效地按数字分组行，并在满足特定条件时递增。我很难弄清楚如何有效地将这些代码转换为可应用于组的代码

以数据帧df为例

df = sqlContext.createDataFrame(
    [
        (33, [], '2017-01-01'),
        (33, ['apple', 'orange'], '2017-01-02'),
        (33, [], '2017-01-03'),
        (33, ['banana'], '2017-01-04')
    ],
    ('ID', 'X', 'date')
)

这段代码实现了我想要的示例df，即按日期排序，并创建组（'grp'），当size列返回到0时，这些组会增加

df \
.withColumn('size', size(col('X'))) \
.withColumn(
    "grp", 
    sum((col('size') == 0).cast("int")).over(Window.orderBy('date'))
).show()

这部分是基于

现在，我要做的是对具有多个ID的数据帧应用相同的方法-实现如下结果

df2 = sqlContext.createDataFrame(
    [
        (33, [], '2017-01-01', 0, 1),
        (33, ['apple', 'orange'], '2017-01-02', 2, 1),
        (33, [], '2017-01-03', 0, 2),
        (33, ['banana'], '2017-01-04', 1, 2),
        (55, ['coffee'], '2017-01-01', 1, 1),
        (55, [], '2017-01-03', 0, 2)
    ],
    ('ID', 'X', 'date', 'size', 'group')
)

为清晰起见进行编辑

1） 对于每个ID的第一个日期-组应为1-无论在任何其他列中显示什么

2） 但是，对于以后的每个日期，我都需要检查“大小”列。如果“大小”列为0，则增加组号。如果它是任何非零的正整数，那么我继续前面的组号

---

我在pandas中看到了一些处理这一问题的方法，但我很难理解pyspark中的应用程序以及pandas与spark中分组数据的不同方式（例如，我是否需要使用称为UADFs的东西？

我添加了一个窗口函数，并在每个ID内创建了一个索引。然后我扩展了条件语句以引用该索引。下面似乎产生了我想要的输出数据帧——但我想知道是否有更有效的方法来实现这一点

window = Window.partitionBy('ID').orderBy('date')
df \
.withColumn('size', size(col('X'))) \
.withColumn('index', rank().over(window).alias('index')) \
.withColumn(
    "grp", 
    sum(((col('size') == 0) | (col('index') == 1)).cast("int")).over(window)
).show()

产生

+---+---------------+----------+----+-----+---+
| ID|              X|      date|size|index|grp|
+---+---------------+----------+----+-----+---+
| 33|             []|2017-01-01|   0|    1|  1|
| 33|[apple, orange]|2017-01-02|   2|    2|  1|
| 33|             []|2017-01-03|   0|    3|  2|
| 33|       [banana]|2017-01-04|   1|    4|  2|
| 55|       [coffee]|2017-01-01|   1|    1|  1|
| 55|             []|2017-01-03|   0|    2|  2|
+---+---------------+----------+----+-----+---+

---

通过检查

大小

是否为零或行是否为第一行，创建列

zero\u或\u first

。然后

求和

df2 = sqlContext.createDataFrame(
    [
        (33, [], '2017-01-01', 0, 1),
        (33, ['apple', 'orange'], '2017-01-02', 2, 1),
        (33, [], '2017-01-03', 0, 2),
        (33, ['banana'], '2017-01-04', 1, 2),
        (55, ['coffee'], '2017-01-01', 1, 1),
        (55, [], '2017-01-03', 0, 2),
        (55, ['banana'], '2017-01-01', 1, 1)
    ],
    ('ID', 'X', 'date', 'size', 'group')
)


w = Window.partitionBy('ID').orderBy('date')
df2 = df2.withColumn('row', F.row_number().over(w))
df2 = df2.withColumn('zero_or_first', F.when((F.col('size')==0)|(F.col('row')==1), 1).otherwise(0))
df2 = df2.withColumn('grp', F.sum('zero_or_first').over(w))
df2.orderBy('ID').show()

这是输出。您可以看到该列

group

=

grp

。其中，

组

是预期结果

+---+---------------+----------+----+-----+---+-------------+---+
| ID|              X|      date|size|group|row|zero_or_first|grp|
+---+---------------+----------+----+-----+---+-------------+---+
| 33|             []|2017-01-01|   0|    1|  1|            1|  1|
| 33|       [banana]|2017-01-04|   1|    2|  4|            0|  2|
| 33|[apple, orange]|2017-01-02|   2|    1|  2|            0|  1|
| 33|             []|2017-01-03|   0|    2|  3|            1|  2|
| 55|       [coffee]|2017-01-01|   1|    1|  1|            1|  1|
| 55|       [banana]|2017-01-01|   1|    1|  2|            0|  1|
| 55|             []|2017-01-03|   0|    2|  3|            1|  2|
+---+---------------+----------+----+-----+---+-------------+---+

---

为什么coffee row的组值为1？它不应该是0吗？@cronoik它是1，因为ID已更改，并且此行是该ID的第一个日期-因此它应该是组1。因此开始值是X的大小，并且每次X的大小为0时它都会递增？开始值始终是1-对于第一个ID+日期。然后，该值仅在X的大小为0时递增。第一个ID+日期的X的大小并不重要（在我的实际数据中，它将始终为0或缺失，但我已经尝试在这里简化）。我将把这些信息编辑到主要问题中，你定义了什么窗口？如果你添加行

（55，['banana']，'2017-01-01'，1，1）

，它将不起作用。问题是我认为应该是

（55，['coffee']，'2017-01-01'，1,0）

。@cronoik为我的答案添加了一个窗口