重置窗口函数的计数器

Question

我有一个如下所示的数据集，我想创建一个新的 C 列，它就像一个计数器/行号，每次 B 列的 0 被 A 的列值分区时都应该重置它

仅使用 SparkSQL / SQL（我可以使用 Pyspark 完成）

>>> rdd = sc.parallelize([
...     [1, 0], [1, 1],[1, 1], [1, 0], [1, 1],
...     [1, 1], [2, 1], [2, 1], [3, 0], [3, 1], [3, 1], [3, 1]])
>>> df = rdd.toDF(['A', 'B'])
>>> 
>>> df.show()
+---+---+
|  A|  B|
+---+---+
|  1|  0|
|  1|  1|
|  1|  1|
|  1|  0|
|  1|  1|
|  1|  1|
|  2|  1|
|  2|  1|
|  3|  0|
|  3|  1|
|  3|  1|
|  3|  1|
+---+---+

我想要达到的目标

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  1|  0|  1|
|  1|  1|  2|
|  1|  1|  3|
|  1|  0|  1|
|  1|  1|  2|
|  1|  1|  3|
|  2|  1|  1|
|  2|  1|  2|
|  3|  0|  1|
|  3|  1|  2|
|  3|  1|  3|
|  3|  1|  4|
+---+---+---+

我目前拥有的东西

>>> spark.sql('''
...    select *, row_number() over(partition by A order by A) as C from df
...    ''').show()

    +---+---+---+
    |  A|  B|  C|
    +---+---+---+
    |  1|  0|  1|
    |  1|  1|  2|
    |  1|  1|  3|
    |  1|  0|  4|
    |  1|  1|  5|
    |  1|  1|  6|
    |  3|  0|  1|
    |  3|  1|  2|
    |  3|  1|  3|
    |  3|  1|  4|
    |  2|  1|  1|
    |  2|  1|  2|
    +---+---+---+

Answer 1

SQL 表代表无序 集合。您需要一个指定数据顺序的列。

使用这样的列，您可以累积 0 值，因为它们似乎是中断。所以：

select df.*, row_number() over (partition by A, grp order by A) as C
from (select df.*,
             sum(case when b = 0 then 1 else 0 end) over (partition by A order by <ordering column>) as grp
      from df
      ) df