向 Pyspark 添加列，将组数分配给相关行

Question

我有一个数据框：

from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('').getOrCreate()
df = spark.createDataFrame([("a", "65"), ("b", "23"),("c", "65"), ("d", "23"),
                        ("a", "66"), ("b", "46"),("c", "23"), ("d", "66"),
            ("b", "5"), ("b", "3"),("c", "3")], ["column2", "value"])
df.show()

+-------+-----+
|column2|value|
+-------+-----+ 
|      a| 65  |
|      b| 23  |
|      c| 65  |
|      d| 23  |
|      a| 66  |
|      b| 46  |
|      c| 23  |
|      d| 66  |
|      b|  5  |
|      b|  3  |
|      c|  3  |
+-------+-----+

我想将每 4 行作为一组。然后关于该组创建新列，我可以将组数分配给相应的行。所以想要的输出如下：

+-------+-----+------+
|column2|value|gr_val|
+-------+-----+ -----+
|      a| 65  |    1 |
|      b| 23  |    1 |
|      c| 65  |    1 |
|      d| 23  |    1 |
|      a| 66  |    2 |
|      b| 46  |    2 |
|      c| 23  |    2 |
|      d| 66  |    2 |
|      b|  5  |    3 |
|      b|  3  |    3 |
|      c|  3  |    3 |
+-------+-----+------+

如果有任何帮助，我将不胜感激！

Answer 1

试试这个方法 -

(1) 创建一个新列（虚拟），该列将按顺序增加每行的数字。 lit('a') 用于创建静态值以生成顺序增加的行号。

(2) 将虚拟列与每组中所需的数字或记录分开（例如 4）并取 ceil。 ceil 返回不小于该值的最小整数。

这里是详细的例子 -

from pyspark.sql.functions import *
from pyspark.sql.window import *

w = Window().partitionBy(lit('a')).orderBy(lit('a'))

df.withColumn("row_num", row_number().over(w))\
    .selectExpr('column2 AS column2','value AS value','ceil(row_num/4) as gr_val')\
    .show()

#+-------+-----+------+
#|column2|value|gr_val|
#+-------+-----+------+
#|      a|   65|     1|
#|      b|   23|     1|
#|      c|   65|     1|
#|      d|   23|     1|
#|      a|   66|     2|
#|      b|   46|     2|
#|      c|   23|     2|
#|      d|   66|     2|
#|      b|    5|     3|
#|      b|    3|     3|
#|      c|    3|     3|
#+-------+-----+------+