Pyspark - 每2行组合2行2个

Question

我在这里有一个pyspark数据框，如下图所示。我想将每两行分组，但是以某种方式：

• 第一行将是来自第1行和第2行的用户
• 第二行将来自第2行和第3行等。

这样的事情：

---CustomerID--previous_stockcodes----stock_codes-----

不使用价格和数量，前一篮子和当前篮子合二为一。例如，CustomerID 12347的第一行将是：

12347----[85116, 22375, 71...]-----[84625A, 84625C, ...]

我已经编写了循环来做到这一点，但这真的效率低，速度慢。我想知道我是否可以使用pyspark有效地做到这一点，但我很难搞清楚。非常感谢提前

Answer 1

您可以使用spark-sql提供的lead函数获取下一行。

lead是一个窗函数。

语法：引导（column_name，int_value，default_value）over（按column_name排序by column_name）

int_value需要您从当前行开始的行数。

default_value在未找到前导行的情况下获取输入

>>> input_df.show()
+----------+---------+----------------+
|customerID|invoiceNo|  stockCode_list|
+----------+---------+----------------+
|     12347|   537626|  [85116, 22375]|
|     12347|   542237|[84625A, 84625C]|
|     12347|   549222|  [22376, 22374]|
|     12347|   556201|  [23084, 23162]|
|     12348|   539318|  [84992, 22951]|
|     12348|   541998|  [21980, 21985]|
|     12348|   548955|  [23077, 23078]|
+----------+---------+----------------+

>>> from pyspark.sql.window import Window
>>> from pyspark.sql.functions import lead,col

>>> win_func = Window.partitionBy("customerID").orderBy("invoiceNo")
>>> new_col = lead("stockCode_list",1,None).over(win_func)
>>> req_df = input_df.select(col("customerID"),col("invoiceNo"),col("stockCode_list"),new_col.alias("req_col"))
>>> req_df.orderBy("customerID","invoiceNo").show()
+----------+---------+----------------+----------------+
|customerID|invoiceNo|  stockCode_list|         req_col|
+----------+---------+----------------+----------------+
|     12347|   537626|  [85116, 22375]|[84625A, 84625C]|
|     12347|   542237|[84625A, 84625C]|  [22376, 22374]|
|     12347|   549222|  [22376, 22374]|  [23084, 23162]|
|     12347|   556201|  [23084, 23162]|            null|
|     12348|   539318|  [84992, 22951]|  [21980, 21985]|
|     12348|   541998|  [21980, 21985]|  [23077, 23078]|
|     12348|   548955|  [23077, 23078]|            null|
+----------+---------+----------------+----------------+