Добавить столбец в 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) Создайте новый столбец (пустышку), в котором будет последовательно увеличиваться число в каждой строке. горит ('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|
#+-------+-----+------+