在pyspark数据框中旋转列和分组的有效方法

时间:2018-06-19 20:01:19

标签: apache-spark pyspark

我在pyspark中有一个数据框,如下所示。

df = spark.createDataFrame([(1,'ios',11,'null'),
                            (1,'ios',12,'null'),
                            (1,'ios',13,'null'),
                            (1,'ios',14,'null'),
                            (1,'android',15,'ok'),
                            (1,'android',16,'not ok'),
                            (1,'android',17,'aborted'),
                            (2,'ios',21,'not ok'),
                            (2,'android',18,'aborted'),
                            (3,'android',18,'null')],
                           ['id','type','s_id','state'])

df.show()
+---+-------+----+-------+
| id|   type|s_id|  state|
+---+-------+----+-------+
|  1|    ios|  11|   null|
|  1|    ios|  12|   null|
|  1|    ios|  13|   null|
|  1|    ios|  14|   null|
|  1|android|  15|     ok|
|  1|android|  16| not_ok|
|  1|android|  17|aborted|
|  2|    ios|  21| not_ok|
|  2|android|  18|aborted|
|  3|android|  18|   null|
+---+-------+----+-------+

现在,我想通过旋转该数据框来创建另一个数据框。
我的操作如下:

from pyspark.sql import Window
from pyspark.sql import functions as f
from pyspark.sql.functions import col, first

windowSpec = Window.partitionBy("id", "type")

df1 = df.withColumn("ranks", f.row_number().over(windowSpec))\
        .filter(f.col("ranks") < 4)\
        .filter(f.col("type") != "")\
        .withColumn("type", f.concat(f.col("type"), 
                    f.col("ranks"))).drop("ranks")\
        .groupBy("id").pivot("type").agg(f.first("s_id"))


df1.show()
+---+--------+--------+--------+----+----+----+
| id|android1|android2|android3|ios1|ios2|ios3|
+---+--------+--------+--------+----+----+----+
|  1|      15|      16|      17|  11|  12|  13|
|  2|      18|    null|    null|  21|null|null|
|  3|      18|    null|    null|null|null|null|
+---+--------+--------+--------+----+----+----+

然后在上述数据框中,我想为每个state加入first valueId
我的操作如下:

过滤state列,其中statenotnull

df2 = df.where(col("state").isNotNull()).groupBy('id').agg(f.first('state').alias('first'))


df2.show()
+---+------+
| id| first|
+---+------+
|  1|    ok|
|  2|not_ok|
+---+------+

加入df1和df2

final_df = df1.join(df2, 'id', 'left_outer')

final_df.show()

+---+--------+--------+--------+----+----+----+------+
| id|android1|android2|android3|ios1|ios2|ios3| first|
+---+--------+--------+--------+----+----+----+------+
|  1|      15|      16|      17|  11|  12|  13|    ok|
|  2|      18|    null|    null|  21|null|null|not_ok|
|  3|      18|    null|    null|null|null|null|  null|
+---+--------+--------+--------+----+----+----+------+

我得到了想要的东西,但想知道是否还有其他有效的方法可以实现这一点。

1 个答案:

答案 0 :(得分:1)

也许,有些效果更有效:

# Compute order of apparition os type
w = Window.partitionBy('id','type').orderBy('s_id')
df = df.withColumn('order',F.rank().over(w))

# Concatenate columns
df = df.withColumn('type',F.concat(F.col('type'),
                                   F.col('order'))).drop('order')
df.show()

+---+--------+----+-------+
| id|    type|s_id|  state|
+---+--------+----+-------+
|  1|    ios1|  11|   null|
|  1|    ios2|  12|   null|
|  1|    ios3|  13|   null|
|  1|    ios4|  14|   null|
|  3|android1|  18|   null|
|  2|    ios1|  21| not ok|
|  2|android1|  18|aborted|
|  1|android1|  15|     ok|
|  1|android2|  16| not ok|
|  1|android3|  17|aborted|
+---+--------+----+-------+

然后旋转数据框,并仅保留前3个os_type列:

# Chose number of cols you want
n_type = 3
l_col=['ios'+str(i+1) for i in range(n_type)]+['android'+str(i+1) for i in range(n_type)]

df = df.groupBy('id').pivot('type').agg({'s_id':'max'}).orderBy('id').select(*l_col)
df.show()

+---+--------+--------+--------+----+----+----+
| id|android1|android2|android3|ios1|ios2|ios3|
+---+--------+--------+--------+----+----+----+
|  1|      15|      16|      17|  11|  12|  13|
|  2|      18|    null|    null|  21|null|null|
|  3|      18|    null|    null|null|null|null|
+---+--------+--------+--------+----+----+----+

然后使用您的方法加入并添加最后一列

编辑::我添加了一个列列表,以仅选择必需的列