Pyspark:如何将行分组为N个组？

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, row_number
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("Grouping Rows").getOrCreate()

df = spark.createDataFrame([(1, "A"), (2, "B"), (3, "C"), (4, "D"), (5, "E"), (6, "F"), (7, "G"), (8, "H")], ["ID", "Value"])

windowSpec = Window.orderBy("ID")
df = df.withColumn("RowNumber", row_number().over(windowSpec))
df.show()

+---+-----+---------+
|ID |Value|RowNumber|
+---+-----+---------+
|1  |A    |1        |
|2  |B    |2        |
|3  |C    |3        |
|4  |D    |4        |
|5  |E    |5        |
|6  |F    |6        |
|7  |G    |7        |
|8  |H    |8        |
+---+-----+---------+

groupSize = 5
df = df.withColumn("Group", ((col("RowNumber")-1) / groupSize).cast("integer"))
df.show()

+---+-----+---------+-----+
|ID |Value|RowNumber|Group|
+---+-----+---------+-----+
|1  |A    |1        |0    |
|2  |B    |2        |0    |
|3  |C    |3        |0    |
|4  |D    |4        |0    |
|5  |E    |5        |0    |
|6  |F    |6        |1    |
|7  |G    |7        |1    |
|8  |H    |8        |1    |
+---+-----+---------+-----+

result = df.groupBy("Group").agg({"ID": "collect_list", "Value": "collect_list"})
result.show()

+-----+------------------+------------------+
|Group|collect_list(ID)  |collect_list(Value)|
+-----+------------------+------------------+
|0    |[1, 2, 3, 4, 5]   |[A, B, C, D, E]   |
|1    |[6, 7, 8]         |[F, G, H]         |
+-----+------------------+------------------+