title: Efficient way to pivot columns and group by in pyspark data frame
body: <p>I have a data frame in pyspark like below.</p>
  
  
  
  <pre><code>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|
  +---+-------+----+-------+
  </code></pre>
  
  <p>Now from this data frame I want to create another data frame by pivoting it.
  <br/>I have done like below :</p>
  
  
  
  <pre><code>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") &lt; 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|
  +---+--------+--------+--------+----+----+----+
  </code></pre>
  
  <p>Then to the above data frame I want to join the <code>state</code> column <code>first value</code> for each <code>Id</code>
  <br/>I have done like below :</p>
  
  <p>Filter <code>state</code> column where <code>state</code> is <code>notnull</code></p>
  
  
  
  <pre><code>df2 = df.where(col("state").isNotNull()).groupBy('id').agg(f.first('state').alias('first'))
  
  
  df2.show()
  +---+------+
  | id| first|
  +---+------+
  |  1|    ok|
  |  2|not_ok|
  +---+------+
  </code></pre>
  
  <p>Join df1 and df2</p>
  
  
  
  <pre><code>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|
  +---+--------+--------+--------+----+----+----+------+
  </code></pre>
  
  <p>I am getting what I want but would like to know if there is any other efficient way to do this.</p>