根據列值的變化對pyspark數據幀進行分區

Question

我在pyspark中有一個數據框。 說有一些列a，b，c ...隨着列的值更改，我想將數據分組。 說

A  B
1  x
1  y
0  x
0  y
0  x
1  y
1  x
1  y

將有3組為(1x,1y),(0x,0y,0x),(1y,1x,1y)和相應的行數據

Answer 1

如果我理解正確，那么您希望每次A列更改值時都創建一個不同的組。

首先，我們將創建一個單調遞增的id，以保持行順序不變：

import pyspark.sql.functions as psf
df = sc.parallelize([[1,'x'],[1,'y'],[0,'x'],[0,'y'],[0,'x'],[1,'y'],[1,'x'],[1,'y']])\
    .toDF(['A', 'B'])\
    .withColumn("rn", psf.monotonically_increasing_id())
df.show()

    +---+---+----------+
    |  A|  B|        rn|
    +---+---+----------+
    |  1|  x|         0|
    |  1|  y|         1|
    |  0|  x|         2|
    |  0|  y|         3|
    |  0|  x|8589934592|
    |  1|  y|8589934593|
    |  1|  x|8589934594|
    |  1|  y|8589934595|
    +---+---+----------+

現在，我們將使用窗口函數創建一個列，每次列A更改時該列包含1 ：

from pyspark.sql import Window
w = Window.orderBy('rn')
df = df.withColumn("changed", (df.A != psf.lag('A', 1, 0).over(w)).cast('int'))

    +---+---+----------+-------+
    |  A|  B|        rn|changed|
    +---+---+----------+-------+
    |  1|  x|         0|      1|
    |  1|  y|         1|      0|
    |  0|  x|         2|      1|
    |  0|  y|         3|      0|
    |  0|  x|8589934592|      0|
    |  1|  y|8589934593|      1|
    |  1|  x|8589934594|      0|
    |  1|  y|8589934595|      0|
    +---+---+----------+-------+

最后，我們將使用另一個窗口函數為每個組分配不同的數字：

df = df.withColumn("group_id", psf.sum("changed").over(w)).drop("rn").drop("changed")

    +---+---+--------+
    |  A|  B|group_id|
    +---+---+--------+
    |  1|  x|       1|
    |  1|  y|       1|
    |  0|  x|       2|
    |  0|  y|       2|
    |  0|  x|       2|
    |  1|  y|       3|
    |  1|  x|       3|
    |  1|  y|       3|
    +---+---+--------+

現在您可以建立小組