如何在Spark数据框中添加具有序列值的列?

时间:2018-08-15 06:07:03

标签: python apache-spark pyspark apache-spark-sql apache-spark-dataset

如何在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

因此,它为所有行增加了一次。如何为每一行增加? 预先感谢。

1 个答案:

答案 0 :(得分:3)

可以使用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)))