PySpark:根据当前行值对行进行计数

时间:2018-06-22 10:03:41

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

我有一个带有“速度”列的DataFrame。
我是否可以有效地为DataFrame中的每一行添加一列,使它们的“速度”在“速度”行的+ / 2范围内?

results = spark.createDataFrame([[1],[2],[3],[4],[5],
                                 [4],[5],[4],[5],[6],
                                 [5],[6],[1],[3],[8],
                                 [2],[5],[6],[10],[12]], 
                                 ['Speed'])

results.show()

+-----+
|Speed|
+-----+
|    1|
|    2|
|    3|
|    4|
|    5|
|    4|
|    5|
|    4|
|    5|
|    6|
|    5|
|    6|
|    1|
|    3|
|    8|
|    2|
|    5|
|    6|
|   10|
|   12|
+-----+

1 个答案:

答案 0 :(得分:1)

您可以使用窗口功能:

# Order the window by speed, and look at range [0;+2]
w = Window.orderBy('Speed').rangeBetween(0,2)

# Define a column counting the number of rows containing value Speed+2
results = results.withColumn('count+2',F.count('Speed').over(w)).orderBy('Speed')
results.show()

+-----+-----+
|Speed|count|
+-----+-----+
|    1|    6|
|    1|    6|
|    2|    7|
|    2|    7|
|    3|   10|
|    3|   10|
|    4|   11|
|    4|   11|
|    4|   11|
|    5|    8|
|    5|    8|
|    5|    8|
|    5|    8|
|    5|    8|
|    6|    4|
|    6|    4|
|    6|    4|
|    8|    2|
|   10|    2|
|   12|    1|
+-----+-----+

注意:窗口函数对所研究的行本身进行计数。您可以通过在计数列中添加-1来纠正此问题

results = results.withColumn('count+2',F.count('Speed').over(w)-1).orderBy('Speed')