Как добавить столбец со значением последовательности в кадре данных Spark?

Question

Как мне добавить столбец со значением последовательности из определенного числа во фрейме данных PySpark?

Текущий набор данных:

Col1    Col2    Flag
Val1    Val2    F
Val3    Val4    T

Но я хочу, чтобы набор данных был таким:

Col1    Col2    Flag    New_Col
Val1    Val2    F       11F
Val3    Val4    T       12T

Я использую код ниже, в Python.

from pyspark.sql import functions as F
from pyspark.sql import types as T

seq = 10

def fn_increment_id(flag):
    global seq
    seq += 1
    return str(seq) + flag

if __name__ == "__main__":
    df = spark.loadFromMapRDB("path/to/table")
    my_udf = F.UserDefinedFunction(fn_increment_id, T.StringType())
    df = df.withColumn("New_Col", my_udf("Flag"))
    print(df.show(10))

Но я заканчиваю с результатом:

Received Dataset:

Col1    Col2    Flag    New_Col
Val1    Val2    F       11F
Val3    Val4    T       11T

Таким образом, он увеличивается один раз для всех строк. Как я могу увеличить для каждой строки? Заранее спасибо.

Answer 1

Столбец с последовательными значениями можно добавить с помощью Window, Это нормально, если датафрейм не слишком большой, для больших фреймов рекомендуется использовать partitionBy в окне, но значения не будут последовательными.

Приведенный ниже код создает порядковые номера для каждой строки, добавляет к ней 10, а затем объединяет значение с Flag столбец для создания нового столбца. Здесь строки отсортированы по Col1 но любой столбец может быть использован.

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

w = Window().orderBy("Col1")
df = df.withColumn("New_Col", concat(row_number().over(w) + 10, col(Flag)))