如何在pyspark中对每个组中的变量进行排序?

时间:2018-05-14 17:43:42

标签: pyspark pyspark-sql

我正在尝试为每个val使用另一列ts对值id进行排序。

# imports
from pyspark.sql import functions as F
from pyspark.sql import SparkSession as ss
import pandas as pd

# create dummy data
pdf = pd.DataFrame( [['2',2,'cat'],['1',1,'dog'],['1',2,'cat'],['2',3,'cat'],['2',4,'dog']] ,columns=['id','ts','val'])
sdf = ss.createDataFrame( pdf )
sdf.show()

+---+---+---+
| id| ts|val|
+---+---+---+
|  2|  2|cat|
|  1|  1|dog|
|  1|  2|cat|
|  2|  3|cat|
|  2|  4|dog|
+---+---+---+

1 个答案:

答案 0 :(得分:1)

您可以按id汇总,然后按ts排序:

sorted_sdf = ( sdf.groupBy('id')
                  .agg( F.sort_array( F.collect_list( F.struct( F.col('ts'), F.col('val') ) ), asc = True)
                  .alias('sorted_col') )  
             )

sorted_sdf.show()

+---+--------------------+
| id|          sorted_col|
+---+--------------------+
|  1|  [[1,dog], [2,cat]]|
|  2|[[2,cat], [3,cat]...|
+---+--------------------+

然后,我们可以爆炸这个列表:

explode_sdf = sorted_sdf.select( 'id' , F.explode( F.col('sorted_col') ).alias('sorted_explode') )

explode_sdf.show()

+---+--------------+
| id|sorted_explode|
+---+--------------+
|  1|       [1,dog]|
|  1|       [2,cat]|
|  2|       [2,cat]|
|  2|       [3,cat]|
|  2|       [4,dog]|
+---+--------------+

sorted_explode的元组分成两部分:

detupled_sdf = explode_sdf.select( 'id', 'sorted_explode.*' )

detupled_sdf.show()

+---+---+---+
| id| ts|val|
+---+---+---+
|  1|  1|dog|
|  1|  2|cat|
|  2|  2|cat|
|  2|  3|cat|
|  2|  4|dog|
+---+---+---+

现在,对于每个ts,我们的原始数据框按id排序!