如何动态地在pyspark数据框中添加列

时间:2018-01-07 05:03:31

标签: pyspark window-functions

我正在尝试根据输入变量vIssueCols

添加几列
from pyspark.sql import HiveContext
from pyspark.sql import functions as F
from pyspark.sql.window import Window
vIssueCols=['jobid','locid']
vQuery1 = 'vSrcData2= vSrcData'
vWindow1 = Window.partitionBy("vKey").orderBy("vOrderBy")
for x in vIssueCols:
Query1=vQuery1+'.withColumn("'+x+'_prev",F.lag(vSrcData.'+x+').over(vWindow1))'

exec(vQuery1)

现在上面的查询将生成如下的vQuery1,它正在工作,但

vSrcData2= vSrcData.withColumn("jobid_prev",F.lag(vSrcData.jobid).over(vWindow1)).withColumn("locid_prev",F.lag(vSrcData.locid).over(vWindow1))

我不能写一个类似

的查询
vSrcData2= vSrcData.withColumn(x+"_prev",F.lag(vSrcData.x).over(vWindow1))for x in vIssueCols

并使用loop语句生成列。有些博客建议添加一个udf并调用它,但是使用udf我将使用上面的执行字符串方法。

1 个答案:

答案 0 :(得分:1)

您可以使用reduce构建查询。

from pyspark.sql.functions import lag
from pyspark.sql.window import Window
from functools import reduce

#sample data
df = sc.parallelize([[1, 200, '1234', 'asdf'],
                     [1, 50, '2345', 'qwerty'],
                     [1, 100, '4567', 'xyz'],
                     [2, 300, '123', 'prem'],
                     [2, 10, '000', 'ankur']]).\
    toDF(["vKey","vOrderBy","jobid","locid"])
df.show()

vWindow1 = Window.partitionBy("vKey").orderBy("vOrderBy")

#your existing processing
df1= df.\
    withColumn("jobid_prev",lag(df.jobid).over(vWindow1)).\
    withColumn("locid_prev",lag(df.locid).over(vWindow1))
df1.show()

#to-be processing
vIssueCols=['jobid','locid']
df2 = (reduce(
    lambda r_df, col_name: r_df.withColumn(col_name+"_prev", lag(r_df[col_name]).over(vWindow1)),
    vIssueCols,
    df
))
df2.show()

示例数据:

+----+--------+-----+------+
|vKey|vOrderBy|jobid| locid|
+----+--------+-----+------+
|   1|     200| 1234|  asdf|
|   1|      50| 2345|qwerty|
|   1|     100| 4567|   xyz|
|   2|     300|  123|  prem|
|   2|      10|  000| ankur|
+----+--------+-----+------+

输出:

+----+--------+-----+------+----------+----------+
|vKey|vOrderBy|jobid| locid|jobid_prev|locid_prev|
+----+--------+-----+------+----------+----------+
|   1|      50| 2345|qwerty|      null|      null|
|   1|     100| 4567|   xyz|      2345|    qwerty|
|   1|     200| 1234|  asdf|      4567|       xyz|
|   2|      10|  000| ankur|      null|      null|
|   2|     300|  123|  prem|       000|     ankur|
+----+--------+-----+------+----------+----------+

希望这有帮助!