Spark Scala SQL 基于另一列创建新列

Question

我有一个 Dataframe 作为

EmpId  EmpName   Salary   SalaryDate 

 1     Amit  1000.0   2016-01-01  

 1     Amit  2000.0   2016-02-01  

 1      Amit     1000.0   2016-03-01   

 1      Amit     2000.0   2016-04-01  

 1      Amit     3000.0   2016-05-01  

 1      Amit     1000.0   2016-06-01

我想添加一个名为 prevSal 的新列，其中包含 Amit 上一行薪水值的数据

预计 Output：

EmpId  EmpName   Salary   SalaryDate prevSal 

 1      Amit     1000.0   2016-01-01  null 

 1      Amit     2000.0   2016-02-01  1000.0 

 1      Amit     1000.0   2016-03-01  2000.0  

 1      Amit     2000.0   2016-04-01  1000.0 

 1      Amit     3000.0   2016-05-01  2000.0 

 1      Amit     1000.0   2016-06-01  3000.0

此外，我想要一个名为 NextSal 的新列，其中包含 Amit 下一行薪水值的数据。 预计 Output

EmpId  EmpName   Salary   SalaryDate prevSal  nextSal 

 1      Amit     1000.0   2016-01-01   null   2000.0  

 1      Amit     2000.0   2016-02-01  1000.0  1000.0 

 1      Amit     1000.0   2016-03-01  2000.0  2000.0 

 1      Amit     2000.0   2016-04-01  1000.0  3000.0 

 1      Amit     3000.0   2016-05-01  2000.0  1000.0 

 1      Amit     1000.0   2016-06-01  3000.0   null 

Answer 1

正如ggordon在他们的评论中指出的那样，这是一个简单的案例，使用lag和lead根据当前每次扫描的行访问上一行和下一行的值。

这些函数起作用的关键基本上是强制执行 DataFrame 中的行顺序，以完全确定哪一行在另一行之前和之后。 为此，我们使用名为orderBy的 Spark window 函数的一个案例。 DataFrame 中的orderBy应该用在将用作对行排序的参考的列上。 由于您的行已经按SalaryDate列排序，因此像这样就足够了：

val w = Window.orderBy("SalaryDate")

关于prevSal列，可以在此处看到lag的用法，根据该信息，我们可以使用上面的 window function 来访问前一行的Salary值，类似这样（第一个参数是要搜索的列的名称，第二个参数是偏移量，也就是我们每次想要向后移动多少行到 go）：

lag("Salary", 1).over(w)

lead的工作方式相同，但用于Salary的远期值。 可以在这里看到lead的用法：

lead("Salary", 1).over(w)

所以所有这些看起来有点像这样（假设df是您的 DataFrame 的名称）：

val w = Window.orderBy("SalaryDate")

df.withColumn("prevSal", lag("Salary", 1).over(w))
  .withColumn("nextSal", lead("Salary", 1).over(w))
  .show()

所以最终结果会是这样的：

+-----+-------+------+----------+-------+-------+
|EmpId|EmpName|Salary|SalaryDate|prevSal|nextSal|
+-----+-------+------+----------+-------+-------+
|    1|   Amit|  1000|2016-01-01|   null|   2000|
|    1|   Amit|  2000|2016-02-01|   1000|   1000|
|    1|   Amit|  1000|2016-03-01|   2000|   2000|
|    1|   Amit|  2000|2016-04-01|   1000|   3000|
|    1|   Amit|  3000|2016-05-01|   2000|   1000|
|    1|   Amit|  1000|2016-06-01|   3000|   null|
+-----+-------+------+----------+-------+-------+

如果您想为每个雇主制作特定的prevSal和nextSal （例如，Amit 有一组不同的prevSal / nextSal比我们说约翰），您只需更改 window function 通过首先按EmpName分区表然后按Salary排序:

val w = Window.partitionBy("EmpName").orderBy("SalaryDate")